Add safetensors reader #555
Open
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
TomNicholas
added
enhancement
TomNicholas
reviewed
Apr 22, 2025
There was a problem hiding this comment.
@nenb this looks great.
Thank you for using the new ManifestStore abstraction too (despite us not having documented or released it yet 😅) - that will save us time refactoring later.
My comments are just nits, seems like we're waiting on zarr dtype extensions to support everything we want here?
| except ImportError: | ||
| raise ImportError("The ml_dtypes package is required to read safetensors files. Please install it with pip install virtualizarr[safetensors].") | ||
|
|
||
| dtype_map = { |
There was a problem hiding this comment.
Maybe just make this a global variable at the top of the file for visibility?
There was a problem hiding this comment.
I'm going to leave it as is for now until the Zarr dtype PR is released, and then will chat with you about how best to implement this function.
nenb
force-pushed
the
add-safetensors-reader
branch
from
e472167 to
b08fc84
Compare
nenb
force-pushed
the
add-safetensors-reader
branch
from
b08fc84 to
6d37b17
Compare
for more information, see https://pre-commit.ci
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
zarr-python.Summary
This PR is the first step in #367. It adds a reader for the official Safetensors specification (see specs). Note that HuggingFace currently implements a SafeTensor reader that goes considerably beyond the official specification. Support for this implementation (which in practice the entire ML/AI community relies on) will be added in the next step for #367.
Notes for reviewers
This PR was developed in association with an AI coding assistant. Make of that what you will!
I have added a notebook to show how everything works. This can be removed in the final version. I have a question at the end of the notebook related to partial get support, which I would appreciate a core dev looking into.
Checklist
docs/releases.rstapi.rst