Skip to main content

zsfm_gguf/
table.rs

1/// Exact-match helper for `tensor_map.rs` files.
2///
3/// Every model crate maps checkpoint names to GGUF names with a first
4/// exact-match block plus custom prefix/layer logic. This helper covers the
5/// exact part so each file keeps one `TABLE` instead of a 15-arm `match`.
6///
7/// Returns `Some(mapped)` on hit, `None` on miss (caller falls through to
8/// prefix logic or returns `None` for skips).
9pub fn map_with_table(name: &str, table: &[(&str, &str)]) -> Option<String> {
10    table
11        .iter()
12        .find(|(from, _)| *from == name)
13        .map(|(_, to)| (*to).to_string())
14}
15
16/// Table-driven exact match with a fallback closure for prefix/layer rules.
17///
18/// Example (TimesFM head):
19/// ```rust
20/// const HEAD: &[(&str, &str)] = &[("a.weight", "b.weight")];
21/// let out = zsfm_gguf::map_with_table("a.weight", HEAD);
22/// assert_eq!(out.as_deref(), Some("b.weight"));
23/// ```
24#[macro_export]
25macro_rules! tensor_map {
26    ($name:expr, { $($from:literal => $to:literal),* $(,)? } $(, $fallback:expr)?) => {{
27        const TABLE: &[(&str, &str)] = &[$(($from, $to)),*];
28        match $crate::map_with_table($name, TABLE) {
29            Some(mapped) => Some(mapped),
30            None => {
31                $( $fallback )?
32                None::<String>
33            }
34        }
35    }};
36}
37
38#[cfg(test)]
39mod tests {
40    use super::*;
41
42    #[test]
43    fn exact_hits() {
44        const T: &[(&str, &str)] = &[("a.weight", "b.weight")];
45        assert_eq!(map_with_table("a.weight", T).as_deref(), Some("b.weight"));
46        assert_eq!(map_with_table("missing", T), None);
47    }
48
49    #[test]
50    fn macro_form() {
51        let out: Option<String> = tensor_map!("a.weight", {"a.weight" => "b.weight"});
52        assert_eq!(out.as_deref(), Some("b.weight"));
53    }
54}