zsfm_ttm/
tensor_map.rs

1/// Map a TinyTimeMixer HuggingFace tensor name to the GGUF naming convention.
2/// Returns `None` for unrecognised names (caller will warn and skip them).
3pub fn map_tensor_name(hf_name: &str) -> Option<String> {
4    // Patcher (Linear patch_length → d_model)
5    match hf_name {
6        "backbone.encoder.patcher.weight" => return Some("enc.patcher.weight".into()),
7        "backbone.encoder.patcher.bias" => return Some("enc.patcher.bias".into()),
8        "decoder.adapter.weight" => return Some("dec.adapter.weight".into()),
9        "decoder.adapter.bias" => return Some("dec.adapter.bias".into()),
10        "head.base_forecast_block.weight" => return Some("head.weight".into()),
11        "head.base_forecast_block.bias" => return Some("head.bias".into()),
12        _ => {}
13    }
14
15    // Encoder adaptive patching mixers:
16    // backbone.encoder.mlp_mixer_encoder.mixers.{L}.mixer_layers.{N}.<part>
17    if let Some(rest) = hf_name.strip_prefix("backbone.encoder.mlp_mixer_encoder.mixers.") {
18        return map_mixer_layer(rest, "enc.blk");
19    }
20
21    // Decoder block mixers:
22    // decoder.decoder_block.mixers.{N}.<part>
23    if let Some(rest) = hf_name.strip_prefix("decoder.decoder_block.mixers.") {
24        return map_decoder_mixer_layer(rest);
25    }
26
27    None
28}
29
30/// Map `{L}.mixer_layers.{N}.<part>` → `{prefix}.{L}.layer.{N}.<gguf_suffix>`
31fn map_mixer_layer(rest: &str, prefix: &str) -> Option<String> {
32    let (level_str, rest) = rest.split_once('.')?;
33    let level: u32 = level_str.parse().ok()?;
34
35    let rest = rest.strip_prefix("mixer_layers.")?;
36    let (layer_str, rest) = rest.split_once('.')?;
37    let layer: u32 = layer_str.parse().ok()?;
38
39    let suffix = map_mixer_suffix(rest)?;
40    Some(format!("{prefix}.{level}.layer.{layer}.{suffix}"))
41}
42
43/// Map `{N}.<part>` → `dec.blk.{N}.<gguf_suffix>`
44fn map_decoder_mixer_layer(rest: &str) -> Option<String> {
45    let (n_str, rest) = rest.split_once('.')?;
46    let n: u32 = n_str.parse().ok()?;
47    let suffix = map_mixer_suffix(rest)?;
48    Some(format!("dec.blk.{n}.{suffix}"))
49}
50
51fn map_mixer_suffix(suffix: &str) -> Option<&'static str> {
52    Some(match suffix {
53        "patch_mixer.norm.norm.weight" => "patch_norm.weight",
54        "patch_mixer.norm.norm.bias" => "patch_norm.bias",
55        "patch_mixer.mlp.fc1.weight" => "patch_fc1.weight",
56        "patch_mixer.mlp.fc1.bias" => "patch_fc1.bias",
57        "patch_mixer.mlp.fc2.weight" => "patch_fc2.weight",
58        "patch_mixer.mlp.fc2.bias" => "patch_fc2.bias",
59        "patch_mixer.gating_block.attn_layer.weight" => "patch_gate.weight",
60        "patch_mixer.gating_block.attn_layer.bias" => "patch_gate.bias",
61        "feature_mixer.norm.norm.weight" => "feat_norm.weight",
62        "feature_mixer.norm.norm.bias" => "feat_norm.bias",
63        "feature_mixer.mlp.fc1.weight" => "feat_fc1.weight",
64        "feature_mixer.mlp.fc1.bias" => "feat_fc1.bias",
65        "feature_mixer.mlp.fc2.weight" => "feat_fc2.weight",
66        "feature_mixer.mlp.fc2.bias" => "feat_fc2.bias",
67        "feature_mixer.gating_block.attn_layer.weight" => "feat_gate.weight",
68        "feature_mixer.gating_block.attn_layer.bias" => "feat_gate.bias",
69        _ => return None,
70    })
71}
72
73#[cfg(test)]
74mod tests {
75    use super::*;
76
77    #[test]
78    fn patcher() {
79        assert_eq!(
80            map_tensor_name("backbone.encoder.patcher.weight"),
81            Some("enc.patcher.weight".into())
82        );
83        assert_eq!(
84            map_tensor_name("backbone.encoder.patcher.bias"),
85            Some("enc.patcher.bias".into())
86        );
87    }
88
89    #[test]
90    fn enc_mixer() {
91        assert_eq!(
92            map_tensor_name("backbone.encoder.mlp_mixer_encoder.mixers.0.mixer_layers.1.patch_mixer.mlp.fc1.weight"),
93            Some("enc.blk.0.layer.1.patch_fc1.weight".into())
94        );
95        assert_eq!(
96            map_tensor_name("backbone.encoder.mlp_mixer_encoder.mixers.2.mixer_layers.0.feature_mixer.gating_block.attn_layer.bias"),
97            Some("enc.blk.2.layer.0.feat_gate.bias".into())
98        );
99    }
100
101    #[test]
102    fn dec_mixer() {
103        assert_eq!(
104            map_tensor_name("decoder.decoder_block.mixers.1.patch_mixer.norm.norm.weight"),
105            Some("dec.blk.1.patch_norm.weight".into())
106        );
107    }
108
109    #[test]
110    fn head() {
111        assert_eq!(
112            map_tensor_name("head.base_forecast_block.weight"),
113            Some("head.weight".into())
114        );
115    }
116
117    #[test]
118    fn unknown() {
119        assert_eq!(map_tensor_name("something.unknown"), None);
120    }
121}
zsfm_ttm/tensor_map.rs

zsfm_ttm/
tensor_map.rs
