Skip to main content

zsfm_timesfm/
tensor_map.rs

1/// Map a HuggingFace TimesFM tensor name to its GGUF blk.N.* equivalent.
2///
3/// HF parameter naming comes from PyTorch nn.Module traversal of
4/// `TimesFM_2p5_200M_torch_module`. Returns `None` for any name that
5/// should not be included in the GGUF (currently none are skipped).
6pub fn map_tensor_name(hf_name: &str) -> Option<String> {
7    const HEAD: &[(&str, &str)] = &[
8        ("tokenizer.hidden_layer.weight", "tokenizer.hidden.weight"),
9        ("tokenizer.hidden_layer.bias", "tokenizer.hidden.bias"),
10        ("tokenizer.output_layer.weight", "tokenizer.output.weight"),
11        ("tokenizer.output_layer.bias", "tokenizer.output.bias"),
12        ("tokenizer.residual_layer.weight", "tokenizer.skip.weight"),
13        ("tokenizer.residual_layer.bias", "tokenizer.skip.bias"),
14        (
15            "output_projection_point.hidden_layer.weight",
16            "out_point.hidden.weight",
17        ),
18        (
19            "output_projection_point.output_layer.weight",
20            "out_point.output.weight",
21        ),
22        (
23            "output_projection_point.residual_layer.weight",
24            "out_point.skip.weight",
25        ),
26        (
27            "output_projection_quantiles.hidden_layer.weight",
28            "out_quantile.hidden.weight",
29        ),
30        (
31            "output_projection_quantiles.output_layer.weight",
32            "out_quantile.output.weight",
33        ),
34        (
35            "output_projection_quantiles.residual_layer.weight",
36            "out_quantile.skip.weight",
37        ),
38    ];
39    if let Some(mapped) = zsfm_gguf::map_with_table(hf_name, HEAD) {
40        return Some(mapped);
41    }
42
43    // Transformer blocks: stacked_xf.{N}.*
44    let rest = hf_name.strip_prefix("stacked_xf.")?;
45    let (block_str, rest) = rest.split_once('.')?;
46    let block: u32 = block_str.parse().ok()?;
47
48    let gguf_suffix = match rest {
49        "pre_attn_ln.scale" => "pre_attn_norm.weight",
50        "post_attn_ln.scale" => "post_attn_norm.weight",
51        "attn.qkv_proj.weight" => "attn_qkv.weight",
52        "attn.out.weight" => "attn_out.weight",
53        "attn.query_ln.scale" => "attn_q_norm.weight",
54        "attn.key_ln.scale" => "attn_k_norm.weight",
55        "attn.per_dim_scale.per_dim_scale" => "attn_q_scale.weight",
56        "pre_ff_ln.scale" => "pre_ff_norm.weight",
57        "post_ff_ln.scale" => "post_ff_norm.weight",
58        "ff0.weight" => "ffn_up.weight",
59        "ff1.weight" => "ffn_down.weight",
60        _ => return None,
61    };
62
63    Some(format!("blk.{block}.{gguf_suffix}"))
64}