zsfm_mitra/
tensor_map.rs

1/// Map Mitra (Tab2D) safetensors tensor names to GGUF naming convention.
2pub fn map_tensor_name(name: &str) -> Option<String> {
3    match name {
4        "final_layer.weight" => return Some("head.weight".into()),
5        "final_layer.bias" => return Some("head.bias".into()),
6        "final_layer_norm.weight" => return Some("norm_f.weight".into()),
7        "final_layer_norm.bias" => return Some("norm_f.bias".into()),
8        "x_embedding.x_embedding.weight" => return Some("x_embed.weight".into()),
9        "x_embedding.x_embedding.bias" => return Some("x_embed.bias".into()),
10        // Classifier: y_embedding.y_embedding is an nn.Embedding (weight only).
11        // Regressor: y_embedding.y_embedding is an nn.Linear (weight + bias).
12        "y_embedding.y_embedding.weight" => return Some("y_embed.weight".into()),
13        "y_embedding.y_embedding.bias" => return Some("y_embed.bias".into()),
14        "y_embedding.y_mask.weight" => return Some("y_mask.weight".into()),
15        _ => {}
16    }
17
18    if let Some(rest) = name.strip_prefix("layers.") {
19        let (n_str, rest) = rest.split_once('.')?;
20        let n: u32 = n_str.parse().ok()?;
21
22        let suffix = match rest {
23            "layer_norm1.weight" => "ln1.weight",
24            "layer_norm1.bias" => "ln1.bias",
25            "layer_norm2.weight" => "ln2.weight",
26            "layer_norm2.bias" => "ln2.bias",
27            "layer_norm3.weight" => "ln3.weight",
28            "layer_norm3.bias" => "ln3.bias",
29            "layer_norm4.weight" => "ln4.weight",
30            "layer_norm4.bias" => "ln4.bias",
31            "attention1.q.weight" => "attn_row.q.weight",
32            "attention1.q.bias" => "attn_row.q.bias",
33            "attention1.k.weight" => "attn_row.k.weight",
34            "attention1.k.bias" => "attn_row.k.bias",
35            "attention1.v.weight" => "attn_row.v.weight",
36            "attention1.v.bias" => "attn_row.v.bias",
37            "attention1.o.weight" => "attn_row.o.weight",
38            "attention1.o.bias" => "attn_row.o.bias",
39            "attention2.q.weight" => "attn_feat.q.weight",
40            "attention2.q.bias" => "attn_feat.q.bias",
41            "attention2.k.weight" => "attn_feat.k.weight",
42            "attention2.k.bias" => "attn_feat.k.bias",
43            "attention2.v.weight" => "attn_feat.v.weight",
44            "attention2.v.bias" => "attn_feat.v.bias",
45            "attention2.o.weight" => "attn_feat.o.weight",
46            "attention2.o.bias" => "attn_feat.o.bias",
47            "linear1.weight" => "mlp1_fc1.weight",
48            "linear1.bias" => "mlp1_fc1.bias",
49            "linear2.weight" => "mlp1_fc2.weight",
50            "linear2.bias" => "mlp1_fc2.bias",
51            "linear3.weight" => "mlp2_fc1.weight",
52            "linear3.bias" => "mlp2_fc1.bias",
53            "linear4.weight" => "mlp2_fc2.weight",
54            "linear4.bias" => "mlp2_fc2.bias",
55            _ => return None,
56        };
57        return Some(format!("blk.{n}.{suffix}"));
58    }
59
60    None
61}
62
63#[cfg(test)]
64mod tests {
65    use super::*;
66
67    #[test]
68    fn top_level() {
69        assert_eq!(map_tensor_name("final_layer.weight"), Some("head.weight".into()));
70        assert_eq!(map_tensor_name("x_embedding.x_embedding.bias"), Some("x_embed.bias".into()));
71        assert_eq!(map_tensor_name("y_embedding.y_mask.weight"), Some("y_mask.weight".into()));
72    }
73
74    #[test]
75    fn block_attn_and_mlp() {
76        assert_eq!(
77            map_tensor_name("layers.0.attention1.q.weight"),
78            Some("blk.0.attn_row.q.weight".into())
79        );
80        assert_eq!(
81            map_tensor_name("layers.11.attention2.o.bias"),
82            Some("blk.11.attn_feat.o.bias".into())
83        );
84        assert_eq!(map_tensor_name("layers.3.linear4.weight"), Some("blk.3.mlp2_fc2.weight".into()));
85    }
86
87    #[test]
88    fn unknown_returns_none() {
89        assert_eq!(map_tensor_name("optimizer.state"), None);
90    }
91}
zsfm_mitra/tensor_map.rs

zsfm_mitra/
tensor_map.rs
