zsfm_moment/
tensor_map.rs

1/// Map MOMENT safetensors tensor names to GGUF naming convention.
2pub fn map_tensor_name(name: &str) -> Option<String> {
3    match name {
4        "patch_embedding.value_embedding.weight"    => return Some("patch_embed.weight".into()),
5        "patch_embedding.position_embedding.pe"     => return Some("pos_embed.pe".into()),
6        "patch_embedding.mask_embedding"            => return Some("mask_embed".into()),
7        "encoder.embed_tokens.weight"               => return Some("token_embed.weight".into()),
8        "encoder.final_layer_norm.weight"           => return Some("norm_f.weight".into()),
9        "head.linear.weight"                        => return Some("head.weight".into()),
10        "head.linear.bias"                          => return Some("head.bias".into()),
11        _ => {}
12    }
13
14    if let Some(rest) = name.strip_prefix("encoder.block.") {
15        let (n_str, rest) = rest.split_once('.')?;
16        let n: u32 = n_str.parse().ok()?;
17
18        let suffix = match rest {
19            "layer.0.SelfAttention.q.weight"                         => "attn_q.weight",
20            "layer.0.SelfAttention.k.weight"                         => "attn_k.weight",
21            "layer.0.SelfAttention.v.weight"                         => "attn_v.weight",
22            "layer.0.SelfAttention.o.weight"                         => "attn_o.weight",
23            "layer.0.SelfAttention.relative_attention_bias.weight"   => "attn_rel_bias.weight",
24            "layer.0.layer_norm.weight"                              => "attn_norm.weight",
25            "layer.1.DenseReluDense.wi_0.weight"                     => "ffn_wi0.weight",
26            "layer.1.DenseReluDense.wi_1.weight"                     => "ffn_wi1.weight",
27            "layer.1.DenseReluDense.wo.weight"                       => "ffn_wo.weight",
28            "layer.1.layer_norm.weight"                              => "ffn_norm.weight",
29            _ => return None,
30        };
31        return Some(format!("blk.{n}.{suffix}"));
32    }
33
34    None
35}
36
37#[cfg(test)]
38mod tests {
39    use super::*;
40
41    #[test]
42    fn patch_embed() {
43        assert_eq!(
44            map_tensor_name("patch_embedding.value_embedding.weight"),
45            Some("patch_embed.weight".into())
46        );
47    }
48
49    #[test]
50    fn block_attn() {
51        assert_eq!(
52            map_tensor_name("encoder.block.0.layer.0.SelfAttention.q.weight"),
53            Some("blk.0.attn_q.weight".into())
54        );
55        assert_eq!(
56            map_tensor_name("encoder.block.23.layer.1.DenseReluDense.wo.weight"),
57            Some("blk.23.ffn_wo.weight".into())
58        );
59    }
60
61    #[test]
62    fn rel_bias_only_in_block0() {
63        assert_eq!(
64            map_tensor_name("encoder.block.0.layer.0.SelfAttention.relative_attention_bias.weight"),
65            Some("blk.0.attn_rel_bias.weight".into())
66        );
67    }
68}
zsfm_moment/tensor_map.rs

zsfm_moment/
tensor_map.rs
