mlmf 0.2.0

Machine Learning Model Files - Loading, saving, and dynamic mapping for ML models
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111

//! SafeTensors format support
//!
//! SafeTensors is the primary format for storing and loading transformer models.
//! This module provides utilities specific to SafeTensors files.

use crate::error::{Error, Result};
use candle_core::{DType, Device, Tensor};
use std::collections::HashMap;
use std::path::Path;

/// Load SafeTensors file with memory mapping
///
/// # Arguments
/// * `file_path` - Path to the .safetensors file
/// * `device` - Target device for tensors
/// * `dtype` - Target data type (tensors will be converted if needed)
///
/// Note: Currently falls back to regular loading due to API differences
pub fn load_mmaped_safetensors<P: AsRef<Path>>(
    file_path: P,
    device: &Device,
    dtype: DType,
) -> Result<HashMap<String, Tensor>> {
    // For now, use regular loading - memory mapping API is different in this Candle version
    load_regular_safetensors(file_path, device, dtype)
}

/// Load SafeTensors file with regular file I/O
///
/// This is safer than memory-mapped loading but potentially slower for large files.
///
/// # Arguments
/// * `file_path` - Path to the .safetensors file
/// * `device` - Target device for tensors
/// * `dtype` - Target data type (tensors will be converted if needed)
pub fn load_regular_safetensors<P: AsRef<Path>>(
    file_path: P,
    device: &Device,
    dtype: DType,
) -> Result<HashMap<String, Tensor>> {
    let file_path = file_path.as_ref();

    let tensors = candle_core::safetensors::load(file_path, device).map_err(|e| {
        Error::model_loading(format!("Failed to load {}: {}", file_path.display(), e))
    })?;

    // Convert to target dtype if needed
    convert_tensors_dtype(tensors, dtype)
}

/// Convert all tensors in a map to the target dtype
fn convert_tensors_dtype(
    tensors: HashMap<String, Tensor>,
    target_dtype: DType,
) -> Result<HashMap<String, Tensor>> {
    let mut converted = HashMap::new();

    for (name, tensor) in tensors {
        let converted_tensor = if tensor.dtype() != target_dtype {
            tensor.to_dtype(target_dtype).map_err(|e| {
                Error::model_loading(format!(
                    "Failed to convert tensor '{}' from {:?} to {:?}: {}",
                    name,
                    tensor.dtype(),
                    target_dtype,
                    e
                ))
            })?
        } else {
            tensor
        };

        converted.insert(name, converted_tensor);
    }

    Ok(converted)
}

/// Get tensor names from SafeTensors file without loading the tensors
///
/// This is useful for inspecting what tensors are available without
/// loading the full model into memory.
pub fn get_safetensors_tensor_names<P: AsRef<Path>>(file_path: P) -> Result<Vec<String>> {
    let file_path = file_path.as_ref();

    // Load just the tensors to get names, then extract keys
    let tensors =
        candle_core::safetensors::load(file_path, &candle_core::Device::Cpu).map_err(|e| {
            Error::model_loading(format!("Failed to load {}: {}", file_path.display(), e))
        })?;

    Ok(tensors.keys().cloned().collect())
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_dtype_conversion() {
        // This test would require actual tensor data
        // For now, just test that the function signature is correct
        let empty_tensors: HashMap<String, Tensor> = HashMap::new();
        let result = convert_tensors_dtype(empty_tensors, DType::F16);
        assert!(result.is_ok());
        assert_eq!(result.unwrap().len(), 0);
    }

    // Additional tests would require actual SafeTensors files
    // These should be in integration tests with real model files
}