tiny-recursive-rs 0.1.0

Rust implementation of Tiny Recursive Models for efficient puzzle solving
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

/// Weight loading from safetensors files
use std::path::Path;
use candle_core::{Device, DType};
use candle_nn::VarBuilder;
use crate::TRMConfig;
use super::TinyRecursiveModel;

/// Load model from safetensors file
///
/// # Arguments
/// * `config` - Model configuration
/// * `weights_path` - Path to safetensors file
/// * `device` - Device to load model on
///
/// # Returns
/// Loaded TinyRecursiveModel
pub fn load_model<P: AsRef<Path>>(
    config: TRMConfig,
    weights_path: P,
    device: &Device,
) -> crate::Result<TinyRecursiveModel> {
    // Load weights using Candle's built-in safetensors support
    let dtype = DType::F32;
    let vb = unsafe {
        VarBuilder::from_mmaped_safetensors(
            &[weights_path.as_ref()],
            dtype,
            device,
        )?
    };

    TinyRecursiveModel::new(config, vb)
}