zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-12-14 17:12:49 +08:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

Metadata support for safetensors (#639)

Browse Source 
* metadata support for safetensors

* aliases making it alittle more readable

* addressing comments

* python binding tests
This commit is contained in:
Diogo
2024-02-08 22:33:15 -05:00
committed by GitHub
parent 221f8d3fc2
commit b57bd0488d
8 changed files with 108 additions and 69 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
python/src/ops.cpp
View File

@@ -3214,8 +3214,9 @@ void init_ops(py::module_& m) {
&mlx_save_safetensor_helper,
"file"_a,
"arrays"_a,
"metadata"_a = none,
R"pbdoc(
save_safetensors(file: str, arrays: Dict[str, array])
save_safetensors(file: str, arrays: Dict[str, array], metadata: Optional[Dict[str, str]] = None)
Save array(s) to a binary file in ``.safetensors`` format.
@@ -3225,6 +3226,7 @@ void init_ops(py::module_& m) {
Args:
file (file, str): File in which the array is saved.
arrays (dict(str, array)): The dictionary of names to arrays to be saved.
metadata (dict(str, str), optional): The dictionary of metadata to be saved.
)pbdoc");
m.def(
"save_gguf",