zsfm_toto/
tensor_map.rs

1/// Convert a Toto HuggingFace tensor name to the GGUF blk.N convention.
2/// Returns `None` for unrecognised names (caller will warn and skip them).
3///
4/// Actual tensor names were confirmed by running `inspect-tensors` on the
5/// downloaded model.safetensors checkpoint.
6pub fn map_tensor_name(hf_name: &str) -> Option<String> {
7    // --- patch projection (two-layer MLP with skip connection) ---
8    match hf_name {
9        "patch_proj.linear1.weight" => return Some("patch_proj.linear1.weight".into()),
10        "patch_proj.linear1.bias"   => return Some("patch_proj.linear1.bias".into()),
11        "patch_proj.linear2.weight" => return Some("patch_proj.linear2.weight".into()),
12        "patch_proj.linear2.bias"   => return Some("patch_proj.linear2.bias".into()),
13        "patch_proj.skip_proj.weight" => return Some("patch_proj.skip_proj.weight".into()),
14        "patch_proj.skip_proj.bias"   => return Some("patch_proj.skip_proj.bias".into()),
15        _ => {}
16    }
17
18    // --- output head projection (two-layer MLP with skip connection) ---
19    match hf_name {
20        "output_head.param_projection.proj.linear1.weight" =>
21            return Some("output_head.linear1.weight".into()),
22        "output_head.param_projection.proj.linear1.bias" =>
23            return Some("output_head.linear1.bias".into()),
24        "output_head.param_projection.proj.linear2.weight" =>
25            return Some("output_head.linear2.weight".into()),
26        "output_head.param_projection.proj.linear2.bias" =>
27            return Some("output_head.linear2.bias".into()),
28        "output_head.param_projection.proj.skip_proj.weight" =>
29            return Some("output_head.skip_proj.weight".into()),
30        "output_head.param_projection.proj.skip_proj.bias" =>
31            return Some("output_head.skip_proj.bias".into()),
32        _ => {}
33    }
34
35    // --- per-block tensors ---
36    // Pattern: transformer.layers.{N}.<suffix>
37    let rest = hf_name.strip_prefix("transformer.layers.")?;
38    let (block_str, suffix) = rest.split_once('.')?;
39    let block: u32 = block_str.parse().ok()?;
40
41    let gguf_suffix = map_block_suffix(suffix)?;
42    Some(format!("blk.{block}.{gguf_suffix}"))
43}
44
45fn map_block_suffix(suffix: &str) -> Option<&'static str> {
46    Some(match suffix {
47        // Fused QKV projection
48        "attn.in_proj.weight"  => "attn_qkv.weight",
49        "attn.in_proj.bias"    => "attn_qkv.bias",
50        // Output projection
51        "attn.out_proj.weight" => "attn_output.weight",
52        "attn.out_proj.bias"   => "attn_output.bias",
53        // Per-dimension scale (Toto-specific learned scaling)
54        "attn._pds.per_dim_scale" => "attn_pds.weight",
55        // Attention temperature (learned scalar per block)
56        "attn_tau" => "attn_tau",
57        // Feed-forward
58        "ffn.fc1.weight" => "ffn_up.weight",
59        "ffn.fc1.bias"   => "ffn_up.bias",
60        "ffn.fc2.weight" => "ffn_down.weight",
61        "ffn.fc2.bias"   => "ffn_down.bias",
62        // MLP temperature (learned scalar per block)
63        "mlp_tau" => "mlp_tau",
64        _ => return None,
65    })
66}
67
68#[cfg(test)]
69mod tests {
70    use super::*;
71
72    #[test]
73    fn patch_proj() {
74        assert_eq!(
75            map_tensor_name("patch_proj.linear1.weight"),
76            Some("patch_proj.linear1.weight".into())
77        );
78        assert_eq!(
79            map_tensor_name("patch_proj.skip_proj.bias"),
80            Some("patch_proj.skip_proj.bias".into())
81        );
82    }
83
84    #[test]
85    fn block_attn_qkv() {
86        assert_eq!(
87            map_tensor_name("transformer.layers.0.attn.in_proj.weight"),
88            Some("blk.0.attn_qkv.weight".into())
89        );
90    }
91
92    #[test]
93    fn block_attn_pds() {
94        assert_eq!(
95            map_tensor_name("transformer.layers.3.attn._pds.per_dim_scale"),
96            Some("blk.3.attn_pds.weight".into())
97        );
98    }
99
100    #[test]
101    fn block_attn_tau() {
102        assert_eq!(
103            map_tensor_name("transformer.layers.47.attn_tau"),
104            Some("blk.47.attn_tau".into())
105        );
106    }
107
108    #[test]
109    fn block_ffn() {
110        assert_eq!(
111            map_tensor_name("transformer.layers.12.ffn.fc1.weight"),
112            Some("blk.12.ffn_up.weight".into())
113        );
114        assert_eq!(
115            map_tensor_name("transformer.layers.12.ffn.fc2.weight"),
116            Some("blk.12.ffn_down.weight".into())
117        );
118    }
119
120    #[test]
121    fn output_head() {
122        assert_eq!(
123            map_tensor_name("output_head.param_projection.proj.linear2.weight"),
124            Some("output_head.linear2.weight".into())
125        );
126        assert_eq!(
127            map_tensor_name("output_head.param_projection.proj.skip_proj.bias"),
128            Some("output_head.skip_proj.bias".into())
129        );
130    }
131
132    #[test]
133    fn unknown_returns_none() {
134        assert_eq!(map_tensor_name("some.unknown.tensor"), None);
135    }
136}
zsfm_toto/tensor_map.rs

zsfm_toto/
tensor_map.rs
