1pub 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
26pub 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}