dnrops
/
candle
mirror of https://github.com/huggingface/candle.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Add a python function to save as safetensors. (#740)

This commit is contained in:
Laurent Mazare 2023-09-04 21:32:14 +02:00 committed by GitHub
parent ab0d9fbdd1
commit 000487c36f
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 13 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
candle-pyo3/src/lib.rs
View File

@ -1,5 +1,4 @@
#![allow(clippy::redundant_closure_call)]
// TODO: Handle negative dimension indexes.
use pyo3::exceptions::{PyTypeError, PyValueError};
use pyo3::prelude::*;
use pyo3::types::{IntoPyDict, PyTuple};
@ -714,6 +713,18 @@ fn load_safetensors(path: &str, py: Python<'_>) -> PyResult<PyObject> {
Ok(res.into_py_dict(py).to_object(py))
}
#[pyfunction]
fn save_safetensors(
path: &str,
tensors: std::collections::HashMap<String, PyTensor>,
) -> PyResult<()> {
let tensors = tensors
.into_iter()
.map(|(s, t)| (s, t.0))
.collect::<std::collections::HashMap<_, _>>();
::candle::safetensors::save(&tensors, path).map_err(wrap_err)
}
#[pyfunction]
fn load_ggml(path: &str, py: Python<'_>) -> PyResult<(PyObject, PyObject, PyObject)> {
let mut file = std::fs::File::open(path)?;
@ -867,6 +878,7 @@ fn candle(py: Python<'_>, m: &PyModule) -> PyResult<()> {
m.add_function(wrap_pyfunction!(rand, m)?)?;
m.add_function(wrap_pyfunction!(randn, m)?)?;
m.add_function(wrap_pyfunction!(tensor, m)?)?;
m.add_function(wrap_pyfunction!(save_safetensors, m)?)?;
m.add_function(wrap_pyfunction!(stack, m)?)?;
m.add_function(wrap_pyfunction!(zeros, m)?)?;
Ok(())