MiniCPM5-1B-rust/src/inference/loader.rs

use std::collections::HashMap;

use crate::model::LlamaForCausalLM;
use burn::module::{Module, Param};
use burn::nn::{EmbeddingRecord, LinearRecord};
use burn::tensor::backend::Backend;
use burn::tensor::{Float, Shape, Tensor, TensorData};
use memmap2::Mmap;
use std::path::Path;

pub fn load_safetensors<B: Backend>(
    model: LlamaForCausalLM<B>,
    path: &Path,
    device: &B::Device,
) -> anyhow::Result<LlamaForCausalLM<B>> {
    let file = std::fs::File::open(path)?;
    let mmap = unsafe { Mmap::map(&file)? };

    let header_len = u64::from_le_bytes(mmap[..8].try_into().unwrap()) as usize;
    let header_bytes = &mmap[8..8 + header_len];
    let header: serde_json::Value = serde_json::from_slice(header_bytes)?;

    let mut tensors = HashMap::new();
    let mut shapes = HashMap::new();
    let mut dtypes = HashMap::new();

    let data_offset = 8 + header_len;

    if let Some(obj) = header.as_object() {
        for (name, info) in obj {
            if name == "__metadata__" {
                continue;
            }
            if let Some(info_obj) = info.as_object() {
                if let Some(offsets) = info_obj.get("data_offsets") {
                    if let Some(arr) = offsets.as_array() {
                        let start = arr[0].as_u64().unwrap() as usize;
                        let end = arr[1].as_u64().unwrap() as usize;
                        let data = &mmap[data_offset + start..data_offset + end];
                        tensors.insert(name.clone(), data.to_vec());
                    }
                }
                if let Some(shape) = info_obj.get("shape") {
                    if let Some(arr) = shape.as_array() {
                        let shape_vec: Vec<usize> = arr
                            .iter()
                            .map(|v| v.as_u64().unwrap() as usize)
                            .collect();
                        shapes.insert(name.clone(), shape_vec);
                    }
                }
                if let Some(dtype) = info_obj.get("dtype") {
                    if let Some(s) = dtype.as_str() {
                        dtypes.insert(name.clone(), s.to_string());
                    }
                }
            }
        }
    }

    let mut model = model;

    if let Some(weight) = load_tensor_2d::<B>(&tensors, &shapes, &dtypes, "model.embed_tokens.weight", device) {
        let record = EmbeddingRecord {
            weight: Param::from_tensor(weight),
        };
        model.model.embed_tokens = model.model.embed_tokens.clone().load_record(record);
    }

    for i in 0..model.config.num_hidden_layers {
        let prefix = format!("model.layers.{i}");

        if let Some(w) = load_tensor_2d::<B>(&tensors, &shapes, &dtypes, &format!("{prefix}.self_attn.q_proj.weight"), device) {
            let record = LinearRecord { weight: Param::from_tensor(w.transpose()), bias: None };
            model.model.layers[i].self_attn.q_proj = model.model.layers[i].self_attn.q_proj.clone().load_record(record);
        }
        if let Some(w) = load_tensor_2d::<B>(&tensors, &shapes, &dtypes, &format!("{prefix}.self_attn.k_proj.weight"), device) {
            let record = LinearRecord { weight: Param::from_tensor(w.transpose()), bias: None };
            model.model.layers[i].self_attn.k_proj = model.model.layers[i].self_attn.k_proj.clone().load_record(record);
        }
        if let Some(w) = load_tensor_2d::<B>(&tensors, &shapes, &dtypes, &format!("{prefix}.self_attn.v_proj.weight"), device) {
            let record = LinearRecord { weight: Param::from_tensor(w.transpose()), bias: None };
            model.model.layers[i].self_attn.v_proj = model.model.layers[i].self_attn.v_proj.clone().load_record(record);
        }
        if let Some(w) = load_tensor_2d::<B>(&tensors, &shapes, &dtypes, &format!("{prefix}.self_attn.o_proj.weight"), device) {
            let record = LinearRecord { weight: Param::from_tensor(w.transpose()), bias: None };
            model.model.layers[i].self_attn.o_proj = model.model.layers[i].self_attn.o_proj.clone().load_record(record);
        }

        if let Some(w) = load_tensor_2d::<B>(&tensors, &shapes, &dtypes, &format!("{prefix}.mlp.gate_proj.weight"), device) {
            let record = LinearRecord { weight: Param::from_tensor(w.transpose()), bias: None };
            model.model.layers[i].mlp.gate_proj = model.model.layers[i].mlp.gate_proj.clone().load_record(record);
        }
        if let Some(w) = load_tensor_2d::<B>(&tensors, &shapes, &dtypes, &format!("{prefix}.mlp.up_proj.weight"), device) {
            let record = LinearRecord { weight: Param::from_tensor(w.transpose()), bias: None };
            model.model.layers[i].mlp.up_proj = model.model.layers[i].mlp.up_proj.clone().load_record(record);
        }
        if let Some(w) = load_tensor_2d::<B>(&tensors, &shapes, &dtypes, &format!("{prefix}.mlp.down_proj.weight"), device) {
            let record = LinearRecord { weight: Param::from_tensor(w.transpose()), bias: None };
            model.model.layers[i].mlp.down_proj = model.model.layers[i].mlp.down_proj.clone().load_record(record);
        }

        if let Some(w) = load_tensor_1d::<B>(&tensors, &shapes, &dtypes, &format!("{prefix}.input_layernorm.weight"), device) {
            model.model.layers[i].input_layernorm.weight = Param::from_tensor(w);
        }
        if let Some(w) = load_tensor_1d::<B>(&tensors, &shapes, &dtypes, &format!("{prefix}.post_attention_layernorm.weight"), device) {
            model.model.layers[i].post_attention_layernorm.weight = Param::from_tensor(w);
        }
    }

    if let Some(w) = load_tensor_1d::<B>(&tensors, &shapes, &dtypes, "model.norm.weight", device) {
        model.model.norm.weight = Param::from_tensor(w);
    }

    if let Some(w) = load_tensor_2d::<B>(&tensors, &shapes, &dtypes, "lm_head.weight", device) {
        let record = LinearRecord { weight: Param::from_tensor(w.transpose()), bias: None };
        model.lm_head = model.lm_head.clone().load_record(record);
    } else if model.config.tie_word_embeddings {
        let embed_weight = model.model.embed_tokens.weight.val();
        let record = LinearRecord { weight: Param::from_tensor(embed_weight.transpose()), bias: None };
        model.lm_head = model.lm_head.clone().load_record(record);
    }

    Ok(model)
}

fn load_tensor_2d<B: Backend>(
    tensors: &HashMap<String, Vec<u8>>,
    shapes: &HashMap<String, Vec<usize>>,
    dtypes: &HashMap<String, String>,
    name: &str,
    device: &B::Device,
) -> Option<Tensor<B, 2, Float>> {
    let data = tensors.get(name)?;
    let shape = shapes.get(name)?;
    let dtype = dtypes.get(name)?;

    assert_eq!(shape.len(), 2, "Expected 2D tensor for {name}");

    let f32_data = convert_to_f32(data, dtype);
    let shape_arr: [usize; 2] = [shape[0], shape[1]];
    let tensor_data = TensorData::new(f32_data, Shape::new(shape_arr));
    Some(Tensor::from_data(tensor_data, device))
}

fn load_tensor_1d<B: Backend>(
    tensors: &HashMap<String, Vec<u8>>,
    shapes: &HashMap<String, Vec<usize>>,
    dtypes: &HashMap<String, String>,
    name: &str,
    device: &B::Device,
) -> Option<Tensor<B, 1, Float>> {
    let data = tensors.get(name)?;
    let shape = shapes.get(name)?;
    let dtype = dtypes.get(name)?;

    assert_eq!(shape.len(), 1, "Expected 1D tensor for {name}");

    let f32_data = convert_to_f32(data, dtype);
    let shape_arr: [usize; 1] = [shape[0]];
    let tensor_data = TensorData::new(f32_data, Shape::new(shape_arr));
    Some(Tensor::from_data(tensor_data, device))
}

fn convert_to_f32(data: &[u8], dtype: &str) -> Vec<f32> {
    match dtype {
        "BF16" => data
            .chunks_exact(2)
            .map(|chunk| {
                let bytes: [u8; 2] = chunk.try_into().unwrap();
                bf16_to_f32(u16::from_le_bytes(bytes))
            })
            .collect(),
        "F16" => data
            .chunks_exact(2)
            .map(|chunk| {
                let bytes: [u8; 2] = chunk.try_into().unwrap();
                f16_to_f32(u16::from_le_bytes(bytes))
            })
            .collect(),
        "F32" => data
            .chunks_exact(4)
            .map(|chunk| {
                let bytes: [u8; 4] = chunk.try_into().unwrap();
                f32::from_le_bytes(bytes)
            })
            .collect(),
        _ => panic!("Unsupported dtype: {dtype}"),
    }
}

fn bf16_to_f32(bf16: u16) -> f32 {
    let f32_bits = (bf16 as u32) << 16;
    f32::from_bits(f32_bits)
}

fn f16_to_f32(f16: u16) -> f32 {
    let sign = (f16 >> 15) & 1;
    let exp = (f16 >> 10) & 0x1F;
    let mant = f16 & 0x3FF;

    if exp == 0 {
        if mant == 0 {
            return if sign == 0 { 0.0 } else { -0.0 };
        }
        let mut m = mant as u32;
        let mut e = 0u32;
        while (m & 0x400) == 0 {
            m <<= 1;
            e += 1;
        }
        m &= 0x3FF;
        let f32_exp = 127 - 14 - e;
        let f32_bits = ((sign as u32) << 31) | (f32_exp << 23) | ((m as u32) << 13);
        return f32::from_bits(f32_bits);
    }

    if exp == 0x1F {
        let f32_bits = ((sign as u32) << 31) | (0xFFu32 << 23) | ((mant as u32) << 13);
        return f32::from_bits(f32_bits);
    }

    let f32_exp = (exp as u32) + 127 - 15;
    let f32_bits = ((sign as u32) << 31) | (f32_exp << 23) | ((mant as u32) << 13);
    f32::from_bits(f32_bits)
}