zsfm_chronos/
tensor_map.rs

1/// Convert a Chronos-2 HuggingFace tensor name to the GGUF naming convention.
2/// Returns `None` for unrecognised names (caller will warn and skip them).
3///
4/// Tensor names are derived from the Python class hierarchy in model.py / layers.py.
5pub fn map_tensor_name(hf_name: &str) -> Option<String> {
6    // Token embedding ([PAD] / [REG])
7    if hf_name == "shared.weight" {
8        return Some("token_embd.weight".into());
9    }
10
11    // Input patch embedding (ResidualBlock: in_dim=3*patch_size, h_dim=d_ff, out_dim=d_model)
12    match hf_name {
13        "input_patch_embedding.hidden_layer.weight"   => return Some("input_patch.hidden.weight".into()),
14        "input_patch_embedding.hidden_layer.bias"     => return Some("input_patch.hidden.bias".into()),
15        "input_patch_embedding.output_layer.weight"   => return Some("input_patch.output.weight".into()),
16        "input_patch_embedding.output_layer.bias"     => return Some("input_patch.output.bias".into()),
17        "input_patch_embedding.residual_layer.weight" => return Some("input_patch.skip.weight".into()),
18        "input_patch_embedding.residual_layer.bias"   => return Some("input_patch.skip.bias".into()),
19        _ => {}
20    }
21
22    // Output patch embedding (ResidualBlock: in_dim=d_model, h_dim=d_ff, out_dim=num_q*patch_size)
23    match hf_name {
24        "output_patch_embedding.hidden_layer.weight"   => return Some("output_patch.hidden.weight".into()),
25        "output_patch_embedding.hidden_layer.bias"     => return Some("output_patch.hidden.bias".into()),
26        "output_patch_embedding.output_layer.weight"   => return Some("output_patch.output.weight".into()),
27        "output_patch_embedding.output_layer.bias"     => return Some("output_patch.output.bias".into()),
28        "output_patch_embedding.residual_layer.weight" => return Some("output_patch.skip.weight".into()),
29        "output_patch_embedding.residual_layer.bias"   => return Some("output_patch.skip.bias".into()),
30        _ => {}
31    }
32
33    // Final encoder layer norm
34    if hf_name == "encoder.final_layer_norm.weight" {
35        return Some("enc_norm.weight".into());
36    }
37
38    // Per-block tensors: encoder.block.{N}.layer.{L}.<suffix>
39    let rest = hf_name.strip_prefix("encoder.block.")?;
40    let (block_str, rest) = rest.split_once('.')?;
41    let block: u32 = block_str.parse().ok()?;
42
43    let rest = rest.strip_prefix("layer.")?;
44    let (layer_str, suffix) = rest.split_once('.')?;
45    let layer: u32 = layer_str.parse().ok()?;
46
47    let gguf_suffix = match layer {
48        0 => map_time_attn_suffix(suffix)?,
49        1 => map_group_attn_suffix(suffix)?,
50        2 => map_ffn_suffix(suffix)?,
51        _ => return None,
52    };
53
54    Some(format!("blk.{block}.{gguf_suffix}"))
55}
56
57fn map_time_attn_suffix(suffix: &str) -> Option<&'static str> {
58    Some(match suffix {
59        "self_attention.q.weight" => "time_attn.q.weight",
60        "self_attention.k.weight" => "time_attn.k.weight",
61        "self_attention.v.weight" => "time_attn.v.weight",
62        "self_attention.o.weight" => "time_attn.o.weight",
63        "layer_norm.weight"       => "time_attn_norm.weight",
64        _ => return None,
65    })
66}
67
68fn map_group_attn_suffix(suffix: &str) -> Option<&'static str> {
69    Some(match suffix {
70        "self_attention.q.weight" => "group_attn.q.weight",
71        "self_attention.k.weight" => "group_attn.k.weight",
72        "self_attention.v.weight" => "group_attn.v.weight",
73        "self_attention.o.weight" => "group_attn.o.weight",
74        "layer_norm.weight"       => "group_attn_norm.weight",
75        _ => return None,
76    })
77}
78
79fn map_ffn_suffix(suffix: &str) -> Option<&'static str> {
80    Some(match suffix {
81        "mlp.wi.weight"    => "ffn.wi.weight",
82        "mlp.wo.weight"    => "ffn.wo.weight",
83        "layer_norm.weight" => "ffn_norm.weight",
84        _ => return None,
85    })
86}
87
88#[cfg(test)]
89mod tests {
90    use super::*;
91
92    #[test]
93    fn token_embd() {
94        assert_eq!(map_tensor_name("shared.weight"), Some("token_embd.weight".into()));
95    }
96
97    #[test]
98    fn input_patch() {
99        assert_eq!(
100            map_tensor_name("input_patch_embedding.hidden_layer.weight"),
101            Some("input_patch.hidden.weight".into())
102        );
103        assert_eq!(
104            map_tensor_name("input_patch_embedding.residual_layer.bias"),
105            Some("input_patch.skip.bias".into())
106        );
107    }
108
109    #[test]
110    fn output_patch() {
111        assert_eq!(
112            map_tensor_name("output_patch_embedding.output_layer.weight"),
113            Some("output_patch.output.weight".into())
114        );
115    }
116
117    #[test]
118    fn enc_norm() {
119        assert_eq!(
120            map_tensor_name("encoder.final_layer_norm.weight"),
121            Some("enc_norm.weight".into())
122        );
123    }
124
125    #[test]
126    fn block_time_attn() {
127        assert_eq!(
128            map_tensor_name("encoder.block.0.layer.0.self_attention.q.weight"),
129            Some("blk.0.time_attn.q.weight".into())
130        );
131        assert_eq!(
132            map_tensor_name("encoder.block.5.layer.0.layer_norm.weight"),
133            Some("blk.5.time_attn_norm.weight".into())
134        );
135    }
136
137    #[test]
138    fn block_group_attn() {
139        assert_eq!(
140            map_tensor_name("encoder.block.3.layer.1.self_attention.o.weight"),
141            Some("blk.3.group_attn.o.weight".into())
142        );
143        assert_eq!(
144            map_tensor_name("encoder.block.3.layer.1.layer_norm.weight"),
145            Some("blk.3.group_attn_norm.weight".into())
146        );
147    }
148
149    #[test]
150    fn block_ffn() {
151        assert_eq!(
152            map_tensor_name("encoder.block.2.layer.2.mlp.wi.weight"),
153            Some("blk.2.ffn.wi.weight".into())
154        );
155        assert_eq!(
156            map_tensor_name("encoder.block.2.layer.2.mlp.wo.weight"),
157            Some("blk.2.ffn.wo.weight".into())
158        );
159        assert_eq!(
160            map_tensor_name("encoder.block.2.layer.2.layer_norm.weight"),
161            Some("blk.2.ffn_norm.weight".into())
162        );
163    }
164
165    #[test]
166    fn unknown_returns_none() {
167        assert_eq!(map_tensor_name("some.unknown.tensor"), None);
168    }
169}
zsfm_chronos/tensor_map.rs

zsfm_chronos/
tensor_map.rs
