zsfm_moirai/
tensor_map.rs

1/// Map Moirai safetensors tensor names to GGUF naming convention.
2pub fn map_tensor_name(name: &str) -> Option<String> {
3    match name {
4        "in_proj.weight" => return Some("in_proj.weight".into()),
5        "in_proj.bias" => return Some("in_proj.bias".into()),
6        "mask_encoding.weight" => return Some("mask_embed.weight".into()),
7        "encoder.norm.weight" => return Some("norm_f.weight".into()),
8        // mixture distribution heads
9        "param_proj.proj.components.0.loc.weight" => return Some("head.st_loc.weight".into()),
10        "param_proj.proj.components.0.loc.bias" => return Some("head.st_loc.bias".into()),
11        "param_proj.proj.components.0.scale.weight" => return Some("head.st_scale.weight".into()),
12        "param_proj.proj.components.0.scale.bias" => return Some("head.st_scale.bias".into()),
13        "param_proj.proj.components.0.df.weight" => return Some("head.st_df.weight".into()),
14        "param_proj.proj.components.0.df.bias" => return Some("head.st_df.bias".into()),
15        "param_proj.proj.weights_logits.weight" => return Some("head.mix_w.weight".into()),
16        "param_proj.proj.weights_logits.bias" => return Some("head.mix_w.bias".into()),
17        // skip unused distribution components (Normal, NB, LogNormal)
18        _ if name.starts_with("param_proj.proj.components.1.")
19            || name.starts_with("param_proj.proj.components.2.")
20            || name.starts_with("param_proj.proj.components.3.") =>
21        {
22            return None
23        }
24        _ => {}
25    }
26
27    if let Some(rest) = name.strip_prefix("encoder.layers.") {
28        let (n_str, rest) = rest.split_once('.')?;
29        let n: u32 = n_str.parse().ok()?;
30
31        let suffix = match rest {
32            "norm1.weight" => "norm1.weight",
33            "norm2.weight" => "norm2.weight",
34            "self_attn.q_proj.weight" => "attn_q.weight",
35            "self_attn.k_proj.weight" => "attn_k.weight",
36            "self_attn.v_proj.weight" => "attn_v.weight",
37            "self_attn.out_proj.weight" => "attn_o.weight",
38            "self_attn.q_norm.weight" => "attn_qn.weight",
39            "self_attn.k_norm.weight" => "attn_kn.weight",
40            "self_attn.var_attn_bias.emb.weight" => "attn_vbias.weight",
41            "ffn.fc1.weight" => "ffn_fc1.weight",
42            "ffn.fc2.weight" => "ffn_fc2.weight",
43            "ffn.fc_gate.weight" => "ffn_gate.weight",
44            _ => return None,
45        };
46        return Some(format!("blk.{n}.{suffix}"));
47    }
48
49    None
50}
51
52#[cfg(test)]
53mod tests {
54    use super::*;
55
56    #[test]
57    fn in_proj() {
58        assert_eq!(
59            map_tensor_name("in_proj.weight"),
60            Some("in_proj.weight".into())
61        );
62    }
63
64    #[test]
65    fn encoder_block() {
66        assert_eq!(
67            map_tensor_name("encoder.layers.0.self_attn.q_proj.weight"),
68            Some("blk.0.attn_q.weight".into())
69        );
70        assert_eq!(
71            map_tensor_name("encoder.layers.23.ffn.fc_gate.weight"),
72            Some("blk.23.ffn_gate.weight".into())
73        );
74    }
75
76    #[test]
77    fn head() {
78        assert_eq!(
79            map_tensor_name("param_proj.proj.components.0.loc.weight"),
80            Some("head.st_loc.weight".into())
81        );
82    }
83}
zsfm_moirai/tensor_map.rs

zsfm_moirai/
tensor_map.rs
