rust-bert to directly support torch.load's format?

[njaard]

In the interests of simplicity, it would be nice of rust-bert directly supported the torch.load format. Is this not already the case simply because torch.load requires python pickle?

If that's the case, then could one not use a crate like serde-picke and just duplicate the deserialization in Rust?

[guillaume-be]

Hello @njaard ,

While not specific to the transformers implementation in this crate, it would indeed be interesting to support loading Python's pickle files directly. I invite you to open an issue on the upstream library tch. or to add to the discussion on the existing https://github.com/LaurentMazare/tch-rs/issues/595 issue

[wallies]

@guillaume-be i could be way off base, but could you look to use some of the work done here https://github.com/rustformers/llm/pull/83. I think they used https://github.com/KerfuffleV2/repugnant-pickle by the looks, which looks to solve some of the issues.

[guillaume-be]

Thank you @wallies for sharing this. With the support for safetensors in tch-rs, and since those are now the standard serialization scheme for the transformers library they are probably a safer and convenient alternative.

Support for direct Pickle loading should probably be handled in the upstream crate (see link above) - although I believe the safetensors support may now be a preferred way forward.

[wallies]

@guillaume-be how would you then load the safetensor file using rust-bert. I've converted using https://huggingface.co/spaces/diffusers/convert

[guillaume-be]

Hello @wallies , The torch bindings will automatically change the internal loading method useds when the file is saved with the safetensors extension: https://github.com/LaurentMazare/tch-rs/blob/e18648b69e53c8a7803e3d5d25381d358202a692/src/nn/var_store.rs#L179

Have you tried loading these weights with this extension?

[wallies]

@guillaume-be i think it was failing cos im just using SentenceEmbeddingsBuilder::local(PATH) and its always trying to find a rust_model.ot file.