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.") => return None,
21        _ => {}
22    }
23
24    if let Some(rest) = name.strip_prefix("encoder.layers.") {
25        let (n_str, rest) = rest.split_once('.')?;
26        let n: u32 = n_str.parse().ok()?;
27
28        let suffix = match rest {
29            "norm1.weight"                         => "norm1.weight",
30            "norm2.weight"                         => "norm2.weight",
31            "self_attn.q_proj.weight"              => "attn_q.weight",
32            "self_attn.k_proj.weight"              => "attn_k.weight",
33            "self_attn.v_proj.weight"              => "attn_v.weight",
34            "self_attn.out_proj.weight"            => "attn_o.weight",
35            "self_attn.q_norm.weight"              => "attn_qn.weight",
36            "self_attn.k_norm.weight"              => "attn_kn.weight",
37            "self_attn.var_attn_bias.emb.weight"   => "attn_vbias.weight",
38            "ffn.fc1.weight"                       => "ffn_fc1.weight",
39            "ffn.fc2.weight"                       => "ffn_fc2.weight",
40            "ffn.fc_gate.weight"                   => "ffn_gate.weight",
41            _ => return None,
42        };
43        return Some(format!("blk.{n}.{suffix}"));
44    }
45
46    None
47}
48
49#[cfg(test)]
50mod tests {
51    use super::*;
52
53    #[test]
54    fn in_proj() {
55        assert_eq!(map_tensor_name("in_proj.weight"), Some("in_proj.weight".into()));
56    }
57
58    #[test]
59    fn encoder_block() {
60        assert_eq!(
61            map_tensor_name("encoder.layers.0.self_attn.q_proj.weight"),
62            Some("blk.0.attn_q.weight".into())
63        );
64        assert_eq!(
65            map_tensor_name("encoder.layers.23.ffn.fc_gate.weight"),
66            Some("blk.23.ffn_gate.weight".into())
67        );
68    }
69
70    #[test]
71    fn head() {
72        assert_eq!(
73            map_tensor_name("param_proj.proj.components.0.loc.weight"),
74            Some("head.st_loc.weight".into())
75        );
76    }
77}
zsfm_moirai/tensor_map.rs

zsfm_moirai/
tensor_map.rs
