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!(
43            map_tensor_name("model.transformer.wte.weight"),
44            Some("enc.wte.weight".into())
45        );
46    }
47
48    #[test]
49    fn block_attn() {
50        assert_eq!(
51            map_tensor_name("model.transformer.h.0.attn.q_proj.weight"),
52            Some("blk.0.attn_q.weight".into())
53        );
54        assert_eq!(
55            map_tensor_name("model.transformer.h.7.mlp.c_proj.weight"),
56            Some("blk.7.mlp_proj.weight".into())
57        );
58    }
59
60    #[test]
61    fn head() {
62        assert_eq!(
63            map_tensor_name("model.param_proj.proj.0.weight"),
64            Some("head.mu.weight".into())
65        );
66    }
67}
zsfm_lag_llama/tensor_map.rs

zsfm_lag_llama/
tensor_map.rs
