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

//! Example: Loading a LLaMA model from SafeTensors
//!
//! This example shows how to load a LLaMA model from a directory containing
//! config.json and .safetensors files using the mlmf library.

use candle_core::{DType, Device};
use mlmf::{loader::load_safetensors, LoadOptions};

fn main() -> Result<(), Box<dyn std::error::Error>> {
    // Configure loading options
    let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu);
    let dtype = match device {
        Device::Cuda(_) => DType::F16, // Use F16 on CUDA for memory efficiency
        _ => DType::F32,               // Use F32 on CPU for compatibility
    };

    println!("🚀 Loading LLaMA model...");
    println!("📱 Device: {:?}", device);
    println!("🔢 Data type: {:?}", dtype);

    let options = LoadOptions::new(device, dtype)
        .with_progress() // Enable progress reporting
        .without_mmap(); // Disable mmap for this example (safer for demo)

    // Load the model - replace with your model path
    let model_path = std::env::args()
        .nth(1)
        .unwrap_or_else(|| "./models/llama-7b".to_string());

    let loaded = load_safetensors(&model_path, options)?;

    // Print model information
    println!("\n✅ Model loaded successfully!");
    println!(
        "🏗️  Architecture: {}",
        loaded
            .name_mapper
            .architecture()
            .map(|arch| arch.name())
            .unwrap_or("Unknown")
    );
    println!("📊 Configuration: {}", loaded.config.summary());
    println!("🧮 Total tensors: {}", loaded.raw_tensors.len());
    println!("🗺️  Mapped tensors: {}", loaded.name_mapper.len());

    // Show some example tensor mappings
    println!("\n📝 Example tensor name mappings:");
    let mut count = 0;
    for (hf_name, mapped_name) in loaded.name_mapper.iter() {
        if count >= 5 {
            println!("   ... and {} more", loaded.name_mapper.len() - count);
            break;
        }
        println!("   {} → {}", hf_name, mapped_name);
        count += 1;
    }

    // Access specific tensors
    println!("\n🔍 Tensor inspection:");
    if let Some(tensor) = loaded.get_tensor("wte.weight") {
        println!(
            "   Token embeddings: {:?} {:?}",
            tensor.dims(),
            tensor.dtype()
        );
    }

    if let Some(tensor) = loaded.get_tensor("h.0.attn.q_proj.weight") {
        println!(
            "   Query projection (layer 0): {:?} {:?}",
            tensor.dims(),
            tensor.dtype()
        );
    }

    // Memory usage information
    let memory_estimate =
        mlmf::validation::estimate_memory_usage(&loaded.config, dtype, Some(1), None);
    println!("\n💾 {}", memory_estimate.summary());

    Ok(())
}