zyx/
module.rs

1// Copyright (C) 2025 zk4x
2// SPDX-License-Identifier: LGPL-3.0-only
3
4use std::{collections::HashMap, ffi::OsStr, fs::File, path::Path};
5
6use crate::{DType, Map, RT, Tensor, ZyxError, shape::Dim};
7
8/// Module trait
9pub trait Module {
10    /// Iterate over all tensors immutably
11    fn iter(&self) -> impl Iterator<Item = &Tensor>;
12
13    /// Iterate over all tensors mutably
14    fn iter_mut(&mut self) -> impl Iterator<Item = &mut Tensor>;
15
16    /// Iterate over tensors without consuming the module
17    fn iter_tensors(&self) -> impl Iterator<Item = (String, &Tensor)>;
18
19    /// From tensors
20    fn iter_tensors_mut(&mut self) -> impl Iterator<Item = (String, &mut Tensor)>;
21
22    /// Realize all tensors in the module
23    ///
24    /// # Errors
25    ///
26    /// Returns error if any tensor cannot be realized.
27    fn realize(&self) -> Result<(), ZyxError> {
28        Tensor::realize(self.iter())
29    }
30
31    /// Set parameters, removes them from params, skips parameters that are not found in params.
32    fn set_params(&mut self, params: &mut HashMap<String, Tensor>) {
33        for (label, tensor) in self.iter_tensors_mut() {
34            if let Some(param) = params.remove(&label) {
35                *tensor = param;
36            }
37        }
38    }
39
40    /// Save tensors or modules to a file determined by file extension.
41    /// Currently only safetensors is supported format.
42    ///
43    /// # Errors
44    ///
45    /// Errors if tensors failed to realize or failed to save to disk.
46    fn save(&self, path: impl AsRef<Path>) -> Result<(), ZyxError> {
47        use std::fmt::Write;
48        use std::io::Write as IOWrite;
49        let mut f = File::create(path)?;
50        let mut header = String::from("{");
51        let mut begin = 0;
52        for (label, tensor) in self.iter_tensors() {
53            let dtype = tensor.dtype();
54            write!(header, "\"{label}\":{{").unwrap();
55            write!(header, "\"dtype\":\"{}\",", dtype.safetensors()).unwrap();
56            let mut st_shape = format!("{:?}", tensor.shape());
57            st_shape.retain(|c| !c.is_whitespace());
58            write!(header, "\"shape\":{st_shape},").unwrap();
59            let size = tensor.numel() * Dim::from(dtype.bit_size() / 8);
60            write!(header, "\"data_offsets\":[{},{}]", begin, begin + size).unwrap();
61            begin += size;
62            write!(header, "}},").unwrap();
63        }
64        header.pop();
65        write!(header, "}}").unwrap();
66        let header_bytes = header.as_bytes();
67        f.write_all(&(header_bytes.len() as u64).to_le_bytes())?;
68        f.write_all(header_bytes)?;
69        for tensor in self.iter() {
70            f.write_all(&tensor.to_le_bytes()?)?;
71        }
72        Ok(())
73    }
74}
75
76/// GGUF metadata
77#[allow(unused)]
78pub enum GGUFMetadataValue {
79    Uint8(u8),
80    Int8(i8),
81    Uint16(u16),
82    Int16(i16),
83    Uint32(u32),
84    Int32(i32),
85    Uint64(u64),
86    Int64(i64),
87    Float32(f32),
88    Float64(f64),
89    Bool(bool),
90    String(String),
91    Array(Box<[GGUFMetadataValue]>),
92}
93
94impl<S: std::hash::BuildHasher + Default> Module for HashMap<String, Tensor, S> {
95    fn iter(&self) -> impl Iterator<Item = &Tensor> {
96        self.values()
97    }
98
99    fn iter_mut(&mut self) -> impl Iterator<Item = &mut Tensor> {
100        self.values_mut()
101    }
102
103    fn iter_tensors(&self) -> impl Iterator<Item = (String, &Tensor)> {
104        self.iter().map(|(k, v): (&String, &Tensor)| (k.clone(), v))
105    }
106
107    fn iter_tensors_mut(&mut self) -> impl Iterator<Item = (String, &mut Tensor)> {
108        self.iter_mut().map(|(k, v): (&String, &mut Tensor)| (k.clone(), v))
109    }
110}
111
112impl Module for Vec<Tensor> {
113    fn iter(&self) -> impl Iterator<Item = &Tensor> {
114        self.into_iter()
115    }
116
117    fn iter_mut(&mut self) -> impl Iterator<Item = &mut Tensor> {
118        self.into_iter()
119    }
120
121    fn iter_tensors(&self) -> impl Iterator<Item = (String, &Tensor)> {
122        self.iter().map(|t: &Tensor| (format!("{}", t.id()), t))
123    }
124
125    fn iter_tensors_mut(&mut self) -> impl Iterator<Item = (String, &mut Tensor)> {
126        self.iter_mut().map(|t: &mut Tensor| (format!("{}", t.id()), t))
127    }
128}
129
130impl<M0: Module, M1: Module> Module for (M0, M1) {
131    fn iter(&self) -> impl Iterator<Item = &Tensor> {
132        self.0.iter().chain(self.1.iter())
133    }
134
135    fn iter_mut(&mut self) -> impl Iterator<Item = &mut Tensor> {
136        self.0.iter_mut().chain(self.1.iter_mut())
137    }
138
139    fn iter_tensors(&self) -> impl Iterator<Item = (String, &Tensor)> {
140        self.0.iter_tensors().chain(self.1.iter_tensors())
141    }
142
143    fn iter_tensors_mut(&mut self) -> impl Iterator<Item = (String, &mut Tensor)> {
144        self.0.iter_tensors_mut().chain(self.1.iter_tensors_mut())
145    }
146}
147
148impl Tensor {
149    /// Load module from path. This function will determine the filetype based on file extension.
150    ///
151    /// # Errors
152    ///
153    /// Errors if loading from disk failed or if loaded tensors could not be allocated to device.
154    #[allow(clippy::missing_panics_doc)]
155    pub fn load(path: impl AsRef<Path>) -> Result<HashMap<String, Tensor>, ZyxError>
156    where
157        Self: Sized,
158    {
159        RT.lock().initialize_devices()?; // So that we load debug mask
160        let e = path.as_ref().extension().and_then(OsStr::to_str).unwrap();
161        match e {
162            "safetensors" => Self::load_safetensors(path),
163            "gguf" => Ok(Self::load_gguf(path)?.1),
164            _ => panic!("Unknown file extension. Zyx currently supports only safetensors format."),
165        }
166    }
167
168    /// Load gguf module from path
169    /// First returned value is metadata, second returned value are named tensors
170    /// # Errors
171    /// read failure
172    #[allow(clippy::missing_panics_doc)]
173    #[allow(clippy::type_complexity)]
174    pub fn load_gguf(path: impl AsRef<Path>) -> Result<(HashMap<String, GGUFMetadataValue>, HashMap<String, Tensor>), ZyxError> {
175        use std::io::Read;
176        let mut f = std::fs::File::open(&path)?;
177        let mut magic = [0; 4];
178        f.read_exact(&mut magic)?;
179        if magic != [b'G', b'G', b'U', b'F'] {
180            if magic == [b'F', b'U', b'G', b'G'] {
181                return Err(ZyxError::parse_error(
182                    "GGUF data seems to be stored in big endian order. Only little endian is supported for GGUF in zyx.".into(),
183                ));
184            }
185            return Err(ZyxError::parse_error(
186                format!("Unknown GGUF magic: {magic:?}. Please check your file.").into(),
187            ));
188        }
189        let mut version_bytes = [0; 4];
190        f.read_exact(&mut version_bytes)?;
191        let version = u32::from_le_bytes(version_bytes);
192        //println!("File size is {} bytes", f.metadata()?.len());
193        let mut tensor_count = [0u8; 8];
194        f.read_exact(&mut tensor_count)?;
195        let tensor_count = u64::from_le_bytes(tensor_count);
196        let mut metadata_kv_count = [0u8; 8];
197        f.read_exact(&mut metadata_kv_count)?;
198        let metadata_kv_count = usize::try_from(u64::from_le_bytes(metadata_kv_count))
199            .map_err(|e| ZyxError::parse_error(format!("Failed to parse tensor count in GGUF file. {e}").into()))?;
200
201        let mut metadata = HashMap::new();
202        for _ in 0..metadata_kv_count {
203            // First string key, (len u64, chars),
204            let mut metadata_key_len = [0; 8];
205            f.read_exact(&mut metadata_key_len)?;
206            let metadata_key_len = u64::from_le_bytes(metadata_key_len);
207            let mut metadata_key_bytes = vec![0u8; usize::try_from(metadata_key_len).unwrap()];
208            f.read_exact(&mut metadata_key_bytes)?;
209            let metadata_key = String::from_utf8(metadata_key_bytes)
210                .map_err(|e| ZyxError::parse_error(format!("GGUF metadata key is not valid UTF-8: {e}").into()))?;
211
212            // Then metadata value type (u32 in GGUF v3, u8 in v1/v2).
213            // Then we the value itself.
214            let metadata_value_type = if version >= 3 {
215                let mut buf = [0; 4];
216                f.read_exact(&mut buf)?;
217                u32::from_le_bytes(buf)
218            } else {
219                let mut buf = [0; 1];
220                f.read_exact(&mut buf)?;
221                u32::from(u8::from_le_bytes(buf))
222            };
223            let metadata_value = match metadata_value_type {
224                0 => {
225                    let mut buf = [0; 1];
226                    f.read_exact(&mut buf)?;
227                    GGUFMetadataValue::Uint8(u8::from_le_bytes(buf))
228                }
229                1 => {
230                    let mut buf = [0; 1];
231                    f.read_exact(&mut buf)?;
232                    GGUFMetadataValue::Int8(i8::from_le_bytes(buf))
233                }
234                2 => {
235                    let mut buf = [0; 2];
236                    f.read_exact(&mut buf)?;
237                    GGUFMetadataValue::Uint16(u16::from_le_bytes(buf))
238                }
239                3 => {
240                    let mut buf = [0; 2];
241                    f.read_exact(&mut buf)?;
242                    GGUFMetadataValue::Int16(i16::from_le_bytes(buf))
243                }
244                4 => {
245                    let mut buf = [0; 4];
246                    f.read_exact(&mut buf)?;
247                    GGUFMetadataValue::Uint32(u32::from_le_bytes(buf))
248                }
249                5 => {
250                    let mut buf = [0; 4];
251                    f.read_exact(&mut buf)?;
252                    GGUFMetadataValue::Int32(i32::from_le_bytes(buf))
253                }
254                6 => {
255                    let mut buf = [0; 4];
256                    f.read_exact(&mut buf)?;
257                    GGUFMetadataValue::Float32(f32::from_le_bytes(buf))
258                }
259                7 => {
260                    let mut buf = [0; 1];
261                    f.read_exact(&mut buf)?;
262                    GGUFMetadataValue::Bool(buf[0] != 0)
263                }
264                8 => {
265                    let mut str_len = [0; 8];
266                    f.read_exact(&mut str_len)?;
267                    let str_len = u64::from_le_bytes(str_len);
268                    let mut s_bytes = vec![0u8; usize::try_from(str_len).unwrap()];
269                    f.read_exact(&mut s_bytes)?;
270                    let s = String::from_utf8(s_bytes)
271                        .map_err(|e| ZyxError::parse_error(format!("GGUF metadata string is not valid UTF-8: {e}").into()))?;
272                    GGUFMetadataValue::String(s)
273                }
274                9 => {
275                    let mut arr_type_buf = [0; 4];
276                    f.read_exact(&mut arr_type_buf)?;
277                    let elem_type = u32::from_le_bytes(arr_type_buf);
278                    let mut arr_len_buf = [0; 8];
279                    f.read_exact(&mut arr_len_buf)?;
280                    let arr_len = u64::from_le_bytes(arr_len_buf);
281                    let mut items = Vec::with_capacity(usize::try_from(arr_len).unwrap());
282                    for _ in 0..arr_len {
283                        let item = match elem_type {
284                            0 => {
285                                let mut buf = [0; 1];
286                                f.read_exact(&mut buf)?;
287                                GGUFMetadataValue::Uint8(u8::from_le_bytes(buf))
288                            }
289                            1 => {
290                                let mut buf = [0; 1];
291                                f.read_exact(&mut buf)?;
292                                GGUFMetadataValue::Int8(i8::from_le_bytes(buf))
293                            }
294                            2 => {
295                                let mut buf = [0; 2];
296                                f.read_exact(&mut buf)?;
297                                GGUFMetadataValue::Uint16(u16::from_le_bytes(buf))
298                            }
299                            3 => {
300                                let mut buf = [0; 2];
301                                f.read_exact(&mut buf)?;
302                                GGUFMetadataValue::Int16(i16::from_le_bytes(buf))
303                            }
304                            4 => {
305                                let mut buf = [0; 4];
306                                f.read_exact(&mut buf)?;
307                                GGUFMetadataValue::Uint32(u32::from_le_bytes(buf))
308                            }
309                            5 => {
310                                let mut buf = [0; 4];
311                                f.read_exact(&mut buf)?;
312                                GGUFMetadataValue::Int32(i32::from_le_bytes(buf))
313                            }
314                            6 => {
315                                let mut buf = [0; 4];
316                                f.read_exact(&mut buf)?;
317                                GGUFMetadataValue::Float32(f32::from_le_bytes(buf))
318                            }
319                            7 => {
320                                let mut buf = [0; 1];
321                                f.read_exact(&mut buf)?;
322                                GGUFMetadataValue::Bool(buf[0] != 0)
323                            }
324                            8 => {
325                                let mut item_len = [0; 8];
326                                f.read_exact(&mut item_len)?;
327                                let item_len = u64::from_le_bytes(item_len);
328                                let mut item_bytes = vec![0u8; usize::try_from(item_len).unwrap()];
329                                f.read_exact(&mut item_bytes)?;
330                                let item = String::from_utf8(item_bytes).map_err(|e| {
331                                    ZyxError::parse_error(format!("GGUF array element string is not valid UTF-8: {e}").into())
332                                })?;
333                                GGUFMetadataValue::String(item)
334                            }
335                            10 => {
336                                let mut buf = [0; 8];
337                                f.read_exact(&mut buf)?;
338                                GGUFMetadataValue::Uint64(u64::from_le_bytes(buf))
339                            }
340                            11 => {
341                                let mut buf = [0; 8];
342                                f.read_exact(&mut buf)?;
343                                GGUFMetadataValue::Int64(i64::from_le_bytes(buf))
344                            }
345                            12 => {
346                                let mut buf = [0; 8];
347                                f.read_exact(&mut buf)?;
348                                GGUFMetadataValue::Float64(f64::from_le_bytes(buf))
349                            }
350                            x => todo!("GGUF array element type {x} not supported"),
351                        };
352                        items.push(item);
353                    }
354                    GGUFMetadataValue::Array(items.into_boxed_slice())
355                }
356                10 => {
357                    let mut buf = [0; 8];
358                    f.read_exact(&mut buf)?;
359                    GGUFMetadataValue::Uint64(u64::from_le_bytes(buf))
360                }
361                11 => {
362                    let mut buf = [0; 8];
363                    f.read_exact(&mut buf)?;
364                    GGUFMetadataValue::Int64(i64::from_le_bytes(buf))
365                }
366                12 => {
367                    let mut buf = [0; 8];
368                    f.read_exact(&mut buf)?;
369                    GGUFMetadataValue::Float64(f64::from_le_bytes(buf))
370                }
371                x => todo!("GGUF metadata type {x} not supported"),
372            };
373            metadata.insert(metadata_key, metadata_value);
374        }
375
376        // First we read the whole description of tensors
377        let mut tensor_header = Map::default();
378        for _ in 0..tensor_count {
379            // name
380            let mut tensor_name_len = [0; 8];
381            f.read_exact(&mut tensor_name_len)?;
382            let tensor_name_len = u64::from_le_bytes(tensor_name_len);
383            let mut tensor_name_bytes = vec![0u8; usize::try_from(tensor_name_len).unwrap()];
384            f.read_exact(&mut tensor_name_bytes)?;
385            let tensor_name = String::from_utf8(tensor_name_bytes)
386                .map_err(|e| ZyxError::parse_error(format!("GGUF tensor name is not valid UTF-8: {e}").into()))?;
387
388            // rank (number of dimensions)
389            let mut rank = [0; 4];
390            f.read_exact(&mut rank)?;
391            let rank = u32::from_le_bytes(rank);
392
393            // shape (NOTE there is no explicit check for endiannes here)
394            let mut shape = vec![0u8; rank as usize * 8];
395            f.read_exact(&mut shape)?;
396            let shape: Vec<Dim> = shape
397                .chunks_exact(8)
398                .map(|x| u64::from_le_bytes([x[0], x[1], x[2], x[3], x[4], x[5], x[6], x[7]]))
399                .collect();
400
401            // dtype
402            let mut dtype = [0; 4];
403            f.read_exact(&mut dtype)?;
404            let dtype = u32::from_le_bytes(dtype);
405            let dtype = match dtype {
406                0 => DType::F32,
407                1 => DType::F16,
408                24 => DType::I8,
409                25 => DType::I16,
410                26 => DType::I32,
411                27 => DType::I64,
412                28 => DType::F64,
413                x => todo!("GGUF dtype {x} is not supported by zyx yet."),
414            };
415
416            // offset (position in file)
417            let mut offset = [0; 8];
418            f.read_exact(&mut offset)?;
419            let offset = u64::from_le_bytes(offset);
420
421            tensor_header.insert(tensor_name, (shape, dtype, offset));
422        }
423
424        let mut progress_bar = if RT.lock().debug.dev() {
425            println!("Loading tensors from safetensors file");
426            let bar = crate::prog_bar::ProgressBar::new(tensor_count);
427            Some(bar)
428        } else {
429            None
430        };
431
432        let mut tensors = HashMap::new();
433        for (name, (shape, dtype, offset)) in tensor_header {
434            if let Some(progress_bar) = &mut progress_bar {
435                progress_bar.inc(1, &format!("{name}, {shape:?}, {dtype}"));
436            }
437            tensors.insert(name, Tensor::from_path(shape, dtype, &path, offset)?);
438        }
439        Ok((metadata, tensors))
440    }
441
442    /// Load safetensors module from path
443    ///
444    /// # Errors
445    /// Errors if path does not exist or IO failed for other reasons.
446    #[allow(clippy::missing_panics_doc)]
447    pub fn load_safetensors(path: impl AsRef<Path>) -> Result<HashMap<String, Tensor>, ZyxError> {
448        use std::io::Read;
449        let mut f = std::fs::File::open(&path)?;
450        //println!("File size is {} bytes", f.metadata()?.len());
451        let mut header_len = [0u8; 8];
452        f.read_exact(&mut header_len)?;
453        let n = usize::try_from(u64::from_le_bytes(header_len))
454            .map_err(|e| ZyxError::parse_error(format!("Failed to parse header len in safetensors file. {e}").into()))?;
455        let mut header = vec![0u8; n];
456        f.read_exact(&mut header)?;
457        let header = core::str::from_utf8(&header).map_err(|err| std::io::Error::new(std::io::ErrorKind::InvalidData, err))?;
458        let mut text = String::with_capacity(10);
459        let mut begin_str = false;
460        let mut i = 0;
461        let mut tensors = HashMap::default();
462        let mut dtype = DType::F32;
463        let mut shape = vec![1];
464        let mut label = String::new();
465        let mut metadata = true;
466        let mut progress_bar = if RT.lock().debug.dev() {
467            println!("Loading tensors from safetensors file");
468            let bar = crate::prog_bar::ProgressBar::new(u64::try_from(header.chars().filter(|&c| c == '[').count()).unwrap() / 2);
469            Some(bar)
470        } else {
471            None
472        };
473        //let mmap = Arc::new(unsafe { memmap2::Mmap::map(&f)? });
474        //let mut mptr = mmap.as_ptr();
475        //mptr = mptr.wrapping_add(8 + header.len());
476        let mut offset = (8 + header.len()) as u64;
477        for x in header.chars() {
478            // We skip metadata for now
479            if metadata && text.starts_with("__metadata__") {
480                if x == '}' {
481                    text.clear();
482                    begin_str = false;
483                    metadata = false;
484                }
485                continue;
486            }
487            if ['"', '[', ']'].contains(&x) {
488                if begin_str {
489                    //std::println!("{text}");
490                    if i % 7 == 0 {
491                        #[allow(clippy::assigning_clones)]
492                        {
493                            label = text.clone();
494                        }
495                    } else if i % 7 == 2 {
496                        dtype = DType::from_safetensors(&text)?;
497                    } else if i % 7 == 4 {
498                        shape = text
499                            .split(',')
500                            .map(|d| {
501                                d.parse::<u64>()
502                                    .map_err(|err| ZyxError::parse_error(format!("Cannot parse safetensors shape: {err}").into()))
503                            })
504                            .collect::<Result<_, ZyxError>>()?;
505                    } else if i % 7 == 6 {
506                        // TODO assert offsets
507                        //println!("Offsets: {text}");
508                        let offsets = text
509                            .split(',')
510                            .map(|offset| {
511                                offset.parse::<u64>().map_err(|err| {
512                                    ZyxError::parse_error(format!("Could not parse safetensors offset: {err}").into())
513                                })
514                            })
515                            .collect::<Result<Vec<_>, ZyxError>>()?;
516                        //println!("Offsets: {offsets:?}");
517                        let bytes = shape.iter().product::<Dim>() * Dim::from(dtype.bit_size() / 8);
518                        if offsets[1] - offsets[0] != bytes {
519                            return Err(ZyxError::parse_error("Safetensors shapes and offsets are incorrect.".into()));
520                        }
521                        if let Some(bar) = &mut progress_bar {
522                            bar.inc(1, &format!("{label}, {shape:?}, {dtype:?}"));
523                        }
524                        let tensor = Tensor::from_path(shape.clone(), dtype, &path, offset)?;
525                        offset += bytes as u64;
526                        tensors.insert(label.clone(), tensor);
527                    }
528                    i += 1;
529                    text.clear();
530                    begin_str = false;
531                } else {
532                    text.clear();
533                    begin_str = true;
534                }
535            } else {
536                text.push(x);
537            }
538        }
539        Ok(tensors)
540    }
541}
zyx/module.rs

zyx/
module.rs
