nuvec_nn/
weights.rs

1//! Loading and saving weights by name ([`crate::Module::load_weights`],
2//! [`crate::Module::save_weights`]).
3
4use std::collections::{BTreeMap, HashMap};
5use std::path::{Path, PathBuf};
6
7use nuvec::tree::{LeafKind, Leaves, LoadMode, Tree, Visit, Visitor};
8use nuvec::{Array, Error, Result};
9
10/// Where [`crate::Module::load_weights`] reads named weights from: a safetensors file (a
11/// path), or named arrays (`Vec<(String, Array)>`, a `HashMap` or `BTreeMap`).
12pub trait WeightSource {
13    /// The named weights.
14    fn weights(self) -> Result<Vec<(String, Array)>>;
15}
16
17fn from_file(path: &Path) -> Result<Vec<(String, Array)>> {
18    Ok(nuvec::io::load_safetensors(path)?
19        .tensors
20        .into_iter()
21        .collect())
22}
23
24impl WeightSource for &Path {
25    fn weights(self) -> Result<Vec<(String, Array)>> {
26        from_file(self)
27    }
28}
29
30impl WeightSource for &PathBuf {
31    fn weights(self) -> Result<Vec<(String, Array)>> {
32        from_file(self)
33    }
34}
35
36impl WeightSource for PathBuf {
37    fn weights(self) -> Result<Vec<(String, Array)>> {
38        from_file(&self)
39    }
40}
41
42impl WeightSource for &str {
43    fn weights(self) -> Result<Vec<(String, Array)>> {
44        from_file(Path::new(self))
45    }
46}
47
48impl WeightSource for String {
49    fn weights(self) -> Result<Vec<(String, Array)>> {
50        from_file(Path::new(&self))
51    }
52}
53
54impl WeightSource for Vec<(String, Array)> {
55    fn weights(self) -> Result<Vec<(String, Array)>> {
56        Ok(self)
57    }
58}
59
60impl<const N: usize> WeightSource for [(String, Array); N] {
61    fn weights(self) -> Result<Vec<(String, Array)>> {
62        Ok(self.into())
63    }
64}
65
66impl<S> WeightSource for HashMap<String, Array, S> {
67    fn weights(self) -> Result<Vec<(String, Array)>> {
68        Ok(self.into_iter().collect())
69    }
70}
71
72impl WeightSource for BTreeMap<String, Array> {
73    fn weights(self) -> Result<Vec<(String, Array)>> {
74        Ok(self.into_iter().collect())
75    }
76}
77
78/// `load_flat`'s error, named for `op`.
79fn renamed(e: Error, op: &str) -> Error {
80    match e {
81        Error::InvalidArgument(m) => Error::InvalidArgument(m.replacen("load_flat", op, 1)),
82        e => e,
83    }
84}
85
86pub(crate) fn load<T: Tree>(tree: &mut T, src: impl WeightSource, mode: LoadMode) -> Result<()> {
87    let op = match mode {
88        LoadMode::Strict => "load_weights",
89        _ => "load_weights_partial",
90    };
91    tree.load_flat(src.weights()?, mode)
92        .map_err(|e| renamed(e, op))
93}
94
95/// Writes the leaves of `tree` (all, or only the trainable ones) as safetensors, named by
96/// path. Leaves that cannot be evaluated (unloaded placeholders) are reported by name.
97pub(crate) fn save<T: Leaves + ?Sized>(tree: &T, path: &Path, trainable_only: bool) -> Result<()> {
98    struct Named(Vec<(String, Array)>, bool);
99    impl Visitor for Named {
100        fn leaf(&mut self, path: &str, a: &Array, _: LeafKind, trainable: bool) {
101            if trainable || !self.1 {
102                self.0.push((path.to_owned(), a.clone()));
103            }
104        }
105    }
106    let op = if trainable_only {
107        "save_trainable"
108    } else {
109        "save_weights"
110    };
111    let mut named = Named(Vec::new(), trainable_only);
112    tree.visit(&mut Visit::new(&mut named));
113    let named = named.0;
114    let mut seen = std::collections::HashSet::new();
115    for (p, _) in &named {
116        if !seen.insert(p.as_str()) {
117            return Err(Error::InvalidArgument(format!(
118                "{op}: the path {p:?} appears twice"
119            )));
120        }
121    }
122    let arrays: Vec<Array> = named.iter().map(|(_, a)| a.clone()).collect();
123    if nuvec::eval(&arrays).is_err() {
124        let failed: Vec<&str> = named
125            .iter()
126            .filter(|(_, a)| a.eval().is_err())
127            .map(|(p, _)| p.as_str())
128            .collect();
129        return Err(Error::InvalidArgument(format!(
130            "{op}: cannot evaluate {} (unloaded placeholders?): {}",
131            if failed.len() == 1 {
132                "a leaf"
133            } else {
134                "leaves"
135            },
136            failed.join(", ")
137        )));
138    }
139    let tensors: Vec<(&str, &Array)> = named.iter().map(|(p, a)| (p.as_str(), a)).collect();
140    nuvec::io::save_safetensors(path, &tensors, &[("format", "mlx")])
141}
nuvec_nn/weights.rs

nuvec_nn/
weights.rs
