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!(map_tensor_name("backbone.encoder.patcher.weight"), Some("enc.patcher.weight".into()));
80        assert_eq!(map_tensor_name("backbone.encoder.patcher.bias"), Some("enc.patcher.bias".into()));
81    }
82
83    #[test]
84    fn enc_mixer() {
85        assert_eq!(
86            map_tensor_name("backbone.encoder.mlp_mixer_encoder.mixers.0.mixer_layers.1.patch_mixer.mlp.fc1.weight"),
87            Some("enc.blk.0.layer.1.patch_fc1.weight".into())
88        );
89        assert_eq!(
90            map_tensor_name("backbone.encoder.mlp_mixer_encoder.mixers.2.mixer_layers.0.feature_mixer.gating_block.attn_layer.bias"),
91            Some("enc.blk.2.layer.0.feat_gate.bias".into())
92        );
93    }
94
95    #[test]
96    fn dec_mixer() {
97        assert_eq!(
98            map_tensor_name("decoder.decoder_block.mixers.1.patch_mixer.norm.norm.weight"),
99            Some("dec.blk.1.patch_norm.weight".into())
100        );
101    }
102
103    #[test]
104    fn head() {
105        assert_eq!(map_tensor_name("head.base_forecast_block.weight"), Some("head.weight".into()));
106    }
107
108    #[test]
109    fn unknown() {
110        assert_eq!(map_tensor_name("something.unknown"), None);
111    }
112}
zsfm_ttm/tensor_map.rs

zsfm_ttm/
tensor_map.rs
