zsfm_lag_llama/
tensor_map.rs

1/// Map a Lag-Llama state_dict tensor name to GGUF naming convention.
2pub fn map_tensor_name(name: &str) -> Option<String> {
3    match name {
4        "model.transformer.wte.weight" => return Some("enc.wte.weight".into()),
5        "model.transformer.wte.bias"   => return Some("enc.wte.bias".into()),
6        "model.transformer.ln_f.scale" => return Some("norm_f.weight".into()),
7        "model.param_proj.proj.0.weight" => return Some("head.mu.weight".into()),
8        "model.param_proj.proj.0.bias"   => return Some("head.mu.bias".into()),
9        "model.param_proj.proj.1.weight" => return Some("head.sigma.weight".into()),
10        "model.param_proj.proj.1.bias"   => return Some("head.sigma.bias".into()),
11        "model.param_proj.proj.2.weight" => return Some("head.nu.weight".into()),
12        "model.param_proj.proj.2.bias"   => return Some("head.nu.bias".into()),
13        _ => {}
14    }
15
16    if let Some(rest) = name.strip_prefix("model.transformer.h.") {
17        let (n_str, rest) = rest.split_once('.')?;
18        let n: u32 = n_str.parse().ok()?;
19        let gguf_suffix = match rest {
20            "rms_1.scale"          => "rms1.weight",
21            "rms_2.scale"          => "rms2.weight",
22            "attn.q_proj.weight"   => "attn_q.weight",
23            "attn.kv_proj.weight"  => "attn_kv.weight",
24            "attn.c_proj.weight"   => "attn_c.weight",
25            "mlp.c_fc1.weight"     => "mlp_fc1.weight",
26            "mlp.c_fc2.weight"     => "mlp_fc2.weight",
27            "mlp.c_proj.weight"    => "mlp_proj.weight",
28            _ => return None,
29        };
30        return Some(format!("blk.{n}.{gguf_suffix}"));
31    }
32
33    None
34}
35
36#[cfg(test)]
37mod tests {
38    use super::*;
39
40    #[test]
41    fn wte() {
42        assert_eq!(map_tensor_name("model.transformer.wte.weight"), Some("enc.wte.weight".into()));
43    }
44
45    #[test]
46    fn block_attn() {
47        assert_eq!(
48            map_tensor_name("model.transformer.h.0.attn.q_proj.weight"),
49            Some("blk.0.attn_q.weight".into())
50        );
51        assert_eq!(
52            map_tensor_name("model.transformer.h.7.mlp.c_proj.weight"),
53            Some("blk.7.mlp_proj.weight".into())
54        );
55    }
56
57    #[test]
58    fn head() {
59        assert_eq!(
60            map_tensor_name("model.param_proj.proj.0.weight"),
61            Some("head.mu.weight".into())
62        );
63    }
64}
zsfm_lag_llama/tensor_map.rs

zsfm_lag_llama/
tensor_map.rs
