Add fs param in both safejax.core.serialize and safejax.core.deserialize
Now safetensors files can be uploaded to any fs defined with fsspec
safetensors files can also be loaded from any fs
When downloading, it also prepares the model params, not just the file download
🔗 Linked Issue/s
19
🧪 Tests
[X] Did you implement unit tests if you need to?
If the above checkbox is checked, could you describe how you unit-tested it?
Still using local filesystem, as still considering whether to run the unit tests in a Docker Container that mocks the cloud storage provider e.g. azurite or moto or just create another process to deploy those servers locally e.g. something similar to what s3fs does.
But there seems to be an issue with that, as mocking does not work with the mock_s3 decorator, nor with the mock_s3 context manager, which would be ideal. Instead, the only thing that seems to work with s3fs is deploying moto locally, and providing that URL to the s3fs instance. In any other case, it fails, even though for other libraries such as boto3 or botocore it works fine...
✨ Features
fsparam in bothsafejax.core.serializeandsafejax.core.deserializesafetensorsfiles can be uploaded to any fs defined withfsspecsafetensorsfiles can also be loaded from any fs🔗 Linked Issue/s
19
🧪 Tests
If the above checkbox is checked, could you describe how you unit-tested it?
Still using local filesystem, as still considering whether to run the unit tests in a Docker Container that mocks the cloud storage provider e.g.
azuriteormotoor just create another process to deploy those servers locally e.g. something similar to whats3fsdoes.But there seems to be an issue with that, as mocking does not work with the
mock_s3decorator, nor with themock_s3context manager, which would be ideal. Instead, the only thing that seems to work withs3fsis deployingmotolocally, and providing that URL to thes3fsinstance. In any other case, it fails, even though for other libraries such asboto3orbotocoreit works fine...