Support importing safetensors format

[wdoppenberg]

Safetensors (de)serialization

Feature description

Load, at either tensor- or module-level, weights from a .safetensors file.

Feature motivation

To be able to tap into the huggingface ecosystem it would be useful to have the ability to (de)serialize for safetensors. This would obviously have the limitation of these files merely containing a key-value JSON for all weights, and no description of the model's topology/architecture just like ONNX does. Ideally this would help grow the Model Zoo.

This would then mostly be a convenience feature, both for sharing and using existing models.

(Optional) Suggest a Solution

Possibly we could start by deriving the TryFrom trait for their TensorView and our Tensor and take it from there. Huggingface's candle (Rust) framework already has the capability to (de)serialize from .safetensors.

[antimora]

I think natively we should stick to the file format we have currently, primarily because it's difficult to maintain many formats. We already support JSON (for visual debugging), Bincode (for embedding data), and MessagePack (for long-term and cross-language support). However, we should support file conversion from other formats, such as safetensors, numpy, and pickle. The conversion code should be placed in burn-import under a feature flag included in the default feature set.

I also would like for us to build a tool to do basic structure conversion. What I mean by that is we need some ability to remap field names and also change the depth/level of field nodes.

[jin-eld]

Hi, is there any timeline for this feature? Given the large number of .safetensor models the lack of a conversion tool does hurt...

[nathanielsimard]

Not yet, but we may implement this sonner rather than later

[antimora]

@nathanielsimard it would be easy to do now that PyTorchFileRecorder works nicely. We just need to do the same things but using safetensors reader which is available from Candle.