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

zsfm_chronos/
tensor_map.rs
