Skip to main content

zsfm_hub/
registry.rs

1/// Canonical HuggingFace repo per friendly model name.
2/// Single source of truth for CLI `--model` defaults and Python `convert()`
3/// dispatch, so the two cannot drift when defaults change.
4pub const MODEL_REPOS: &[(&str, &str)] = &[
5    ("toto", "Datadog/Toto-2.0-2.5B"),
6    ("chronos", "amazon/chronos-2"),
7    ("timesfm", "google/timesfm-2.5-200m-pytorch"),
8    ("sundial", "thuml/sundial-base-128m"),
9    ("ttm", "ibm-granite/granite-timeseries-ttm-r2"),
10    ("lag_llama", "time-series-foundation-models/Lag-Llama"),
11    ("lag-llama", "time-series-foundation-models/Lag-Llama"),
12    ("moment", "AutonLab/MOMENT-1-large"),
13    ("moirai", "Salesforce/moirai-1.0-R-large"),
14    ("moirai2", "Salesforce/moirai-2.0-R-small"),
15    ("moirai-2", "Salesforce/moirai-2.0-R-small"),
16    ("flowstate", "ibm-granite/granite-timeseries-flowstate-r1"),
17    ("tirex", "NX-AI/TiRex"),
18    ("mitra-classifier", "autogluon/mitra-classifier"),
19    ("mitra-regression", "autogluon/mitra-regressor"),
20    ("tabdpt", "Layer6/TabDPT"),
21    ("tabicl", "jingang/TabICL"),
22    ("tabpfn", "Prior-Labs/tabpfn_3"),
23    ("tabfm", "google/tabfm-1.0.0-pytorch"),
24];
25
26/// Look up the default repo for a friendly model name.
27pub fn default_repo_for(model: &str) -> Option<&'static str> {
28    MODEL_REPOS
29        .iter()
30        .find(|(name, _)| *name == model)
31        .map(|(_, repo)| *repo)
32}
33
34#[cfg(test)]
35mod tests {
36    use super::*;
37
38    #[test]
39    fn every_model_has_repo() {
40        for m in [
41            "toto",
42            "chronos",
43            "timesfm",
44            "sundial",
45            "ttm",
46            "moment",
47            "moirai",
48            "moirai2",
49            "flowstate",
50            "tirex",
51            "tabdpt",
52            "tabicl",
53            "tabpfn",
54            "tabfm",
55        ] {
56            assert!(default_repo_for(m).is_some(), "missing repo for {m}");
57        }
58    }
59}