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        }
23        "output_head.param_projection.proj.linear1.bias" => {
24            return Some("output_head.linear1.bias".into())
25        }
26        "output_head.param_projection.proj.linear2.weight" => {
27            return Some("output_head.linear2.weight".into())
28        }
29        "output_head.param_projection.proj.linear2.bias" => {
30            return Some("output_head.linear2.bias".into())
31        }
32        "output_head.param_projection.proj.skip_proj.weight" => {
33            return Some("output_head.skip_proj.weight".into())
34        }
35        "output_head.param_projection.proj.skip_proj.bias" => {
36            return Some("output_head.skip_proj.bias".into())
37        }
38        _ => {}
39    }
40
41    // --- per-block tensors ---
42    // Pattern: transformer.layers.{N}.<suffix>
43    let rest = hf_name.strip_prefix("transformer.layers.")?;
44    let (block_str, suffix) = rest.split_once('.')?;
45    let block: u32 = block_str.parse().ok()?;
46
47    let gguf_suffix = map_block_suffix(suffix)?;
48    Some(format!("blk.{block}.{gguf_suffix}"))
49}
50
51fn map_block_suffix(suffix: &str) -> Option<&'static str> {
52    Some(match suffix {
53        // Fused QKV projection
54        "attn.in_proj.weight" => "attn_qkv.weight",
55        "attn.in_proj.bias" => "attn_qkv.bias",
56        // Output projection
57        "attn.out_proj.weight" => "attn_output.weight",
58        "attn.out_proj.bias" => "attn_output.bias",
59        // Per-dimension scale (Toto-specific learned scaling)
60        "attn._pds.per_dim_scale" => "attn_pds.weight",
61        // Attention temperature (learned scalar per block)
62        "attn_tau" => "attn_tau",
63        // Feed-forward
64        "ffn.fc1.weight" => "ffn_up.weight",
65        "ffn.fc1.bias" => "ffn_up.bias",
66        "ffn.fc2.weight" => "ffn_down.weight",
67        "ffn.fc2.bias" => "ffn_down.bias",
68        // MLP temperature (learned scalar per block)
69        "mlp_tau" => "mlp_tau",
70        _ => return None,
71    })
72}
73
74#[cfg(test)]
75mod tests {
76    use super::*;
77
78    #[test]
79    fn patch_proj() {
80        assert_eq!(
81            map_tensor_name("patch_proj.linear1.weight"),
82            Some("patch_proj.linear1.weight".into())
83        );
84        assert_eq!(
85            map_tensor_name("patch_proj.skip_proj.bias"),
86            Some("patch_proj.skip_proj.bias".into())
87        );
88    }
89
90    #[test]
91    fn block_attn_qkv() {
92        assert_eq!(
93            map_tensor_name("transformer.layers.0.attn.in_proj.weight"),
94            Some("blk.0.attn_qkv.weight".into())
95        );
96    }
97
98    #[test]
99    fn block_attn_pds() {
100        assert_eq!(
101            map_tensor_name("transformer.layers.3.attn._pds.per_dim_scale"),
102            Some("blk.3.attn_pds.weight".into())
103        );
104    }
105
106    #[test]
107    fn block_attn_tau() {
108        assert_eq!(
109            map_tensor_name("transformer.layers.47.attn_tau"),
110            Some("blk.47.attn_tau".into())
111        );
112    }
113
114    #[test]
115    fn block_ffn() {
116        assert_eq!(
117            map_tensor_name("transformer.layers.12.ffn.fc1.weight"),
118            Some("blk.12.ffn_up.weight".into())
119        );
120        assert_eq!(
121            map_tensor_name("transformer.layers.12.ffn.fc2.weight"),
122            Some("blk.12.ffn_down.weight".into())
123        );
124    }
125
126    #[test]
127    fn output_head() {
128        assert_eq!(
129            map_tensor_name("output_head.param_projection.proj.linear2.weight"),
130            Some("output_head.linear2.weight".into())
131        );
132        assert_eq!(
133            map_tensor_name("output_head.param_projection.proj.skip_proj.bias"),
134            Some("output_head.skip_proj.bias".into())
135        );
136    }
137
138    #[test]
139    fn unknown_returns_none() {
140        assert_eq!(map_tensor_name("some.unknown.tensor"), None);
141    }
142}
zsfm_toto/tensor_map.rs

zsfm_toto/
tensor_map.rs
