kime_engine/
hub.rs

1//! Finding a model by name, in the Hugging Face cache layout so downloads made by Laya or
2//! `huggingface_hub` are reused and ours are reused by them.
3//!
4//! A name is a local path (a checkpoint directory or a `.kime` file), `hf://org/repo` with an
5//! optional subfolder after it, or one of the aliases in [`ALIASES`]. Nothing here touches the
6//! network. `kime pull` does the downloading, into the same layout.
7
8use std::path::{Path, PathBuf};
9
10/// The published compat models: alias, repo and subfolder.
11pub const ALIASES: [(&str, &str, &str); 2] = [
12    ("laya", "convaiinnovations/laya", ""),
13    ("laya-multilingual", "convaiinnovations/laya", "multilingual"),
14];
15
16/// The files a compat checkpoint needs, relative to its folder in the repo.
17pub const FILES: [&str; 5] = [
18    "rl_agent_config.json",
19    "encoder/config.json",
20    "tokenizer/tokenizer.json",
21    "tokenizer/tokenizer_config.json",
22    "model.safetensors",
23];
24
25/// Where a name points on the hub.
26#[derive(Debug, Clone, PartialEq, Eq)]
27pub struct HubRef {
28    /// `org/repo`.
29    pub repo: String,
30    /// The folder inside the repo, empty for the root.
31    pub subfolder: String,
32}
33
34impl HubRef {
35    /// Reads an alias or an `hf://org/repo[/subfolder]` reference.
36    #[must_use]
37    pub fn parse(name: &str) -> Option<HubRef> {
38        if let Some((_, repo, sub)) = ALIASES.iter().find(|a| a.0 == name) {
39            return Some(HubRef { repo: (*repo).into(), subfolder: (*sub).into() });
40        }
41        let rest = name.strip_prefix("hf://")?;
42        let mut parts = rest.splitn(3, '/');
43        let (org, repo) = (parts.next()?, parts.next()?);
44        if org.is_empty() || repo.is_empty() {
45            return None;
46        }
47        let subfolder = parts.next().unwrap_or("").trim_matches('/').to_string();
48        Some(HubRef { repo: format!("{org}/{repo}"), subfolder })
49    }
50
51    /// The repo's folder in the cache.
52    #[must_use]
53    pub fn repo_dir(&self, cache: &Path) -> PathBuf {
54        cache.join(format!("models--{}", self.repo.replace('/', "--")))
55    }
56
57    /// The checkpoint folder of the snapshot `refs/main` points at, if there is one.
58    #[must_use]
59    pub fn local(&self, cache: &Path) -> Option<PathBuf> {
60        let repo = self.repo_dir(cache);
61        let rev = std::fs::read_to_string(repo.join("refs/main")).ok()?;
62        let dir = repo.join("snapshots").join(rev.trim()).join(&self.subfolder);
63        dir.join("model.safetensors").is_file().then_some(dir)
64    }
65}
66
67/// The hub cache: `$HF_HUB_CACHE`, else `$HF_HOME/hub`, else `~/.cache/huggingface/hub`.
68#[must_use]
69pub fn cache_dir() -> PathBuf {
70    let var = |k: &str| std::env::var_os(k).filter(|v| !v.is_empty()).map(PathBuf::from);
71    if let Some(d) = var("HF_HUB_CACHE") {
72        return d;
73    }
74    if let Some(d) = var("HF_HOME") {
75        return d.join("hub");
76    }
77    let home = var("HOME").or_else(|| var("USERPROFILE")).unwrap_or_else(|| PathBuf::from("."));
78    home.join(".cache/huggingface/hub")
79}
80
81/// The path a model name refers to on this machine.
82///
83/// # Errors
84///
85/// A message saying what was looked for and how to fetch it.
86pub fn resolve(name: &str) -> Result<PathBuf, String> {
87    let path = Path::new(name);
88    if path.exists() {
89        return Ok(path.to_path_buf());
90    }
91    let Some(r) = HubRef::parse(name) else {
92        let known: Vec<&str> = ALIASES.iter().map(|a| a.0).collect();
93        return Err(format!(
94            "no model {name:?}: not a path, an hf:// reference or one of {}",
95            known.join(", ")
96        ));
97    };
98    let cache = cache_dir();
99    r.local(&cache)
100        .ok_or_else(|| format!("{name} is not in {}, run kime pull {name} first", cache.display()))
101}
102
103#[cfg(test)]
104mod tests {
105    use super::*;
106
107    #[test]
108    fn names() {
109        let r = HubRef::parse("laya-multilingual").unwrap();
110        assert_eq!(
111            (r.repo.as_str(), r.subfolder.as_str()),
112            ("convaiinnovations/laya", "multilingual")
113        );
114        let r = HubRef::parse("hf://org/repo/a/b/").unwrap();
115        assert_eq!((r.repo.as_str(), r.subfolder.as_str()), ("org/repo", "a/b"));
116        assert_eq!(HubRef::parse("hf://org"), None);
117        assert_eq!(HubRef::parse("kime-v1-s-en"), None);
118        assert_eq!(r.repo_dir(Path::new("/c")), Path::new("/c/models--org--repo"));
119    }
120}
kime_engine/hub.rs

kime_engine/
hub.rs
