Support loading from model.safetensors.index.json

[findmyway]

No description provided.

[pxl-th]

LGTM! Would be good to add a tests if possible.

[ToucheSir]

I noticed that this index file is not covered anywhere in https://huggingface.co/docs/safetensors/index or the safetensors repo. Is it a huggingface-specific thing? Is there a documentation link we can point users to for it?

[findmyway]

I noticed that this index file is not covered anywhere in https://huggingface.co/docs/safetensors/index or the safetensors repo. Is it a huggingface-specific thing? Is there a documentation link we can point users to for it?

See https://huggingface.co/docs/huggingface_hub/package_reference/hf_api#huggingface_hub.utils.SafetensorsRepoMetadata

[findmyway]

LGTM! Would be good to add a tests if possible.

Added

[chengchingwen]

The index.json is defined by huggingface in their python binding for sharding the model weights and is not part of the safetensors format/spec. This should not be in this package (or at least, it should be named as something like load_shard_safetensors instead of modifying load_safetensors.

[findmyway]

Hi @chengchingwen ,

I agree this feature is not part of the spec. The reason is that in real world cases, different packages may have different approaches to handle the shards (like loading them distributedly or GC during loading). However, I think the modifications I made here provide a nice-to-have fallback (by loading them all in memory).

Regarding the naming issue, reusing the load_safetensors won't create any breaking change. But I'd be happy to change if you insist.

[chengchingwen]

The reason is that in real world cases, different packages may have different approaches to handle the shards

It's also part of the reason that I think it should not be in this package, but I agree it would be convenient to have and is a reasonable default for loading sharded weights. Personally, the index.json is kind of unsafe regarding the original intention of safetensors. Meanwhile, it doesn't have a stable format/spec which make things break easily and inconspicuously if any changes have been made to their python code. Adding new function also won't be a breaking change and could provide a clear separation between spec and custom behavior, so I would insist the name change if we are going to merge it.

[chengchingwen]

Also, some unneeded files should be removed. e.g. test/README.md