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); 3] = [
12    ("laya", "convaiinnovations/laya", ""),
13    ("laya-multilingual", "convaiinnovations/laya", "multilingual"),
14    ("laya-typed-decisions", "convaiinnovations/laya", "typed-decisions"),
15];
16
17/// The files a compat checkpoint needs, relative to its folder in the repo.
18pub const FILES: [&str; 5] = [
19    "rl_agent_config.json",
20    "encoder/config.json",
21    "tokenizer/tokenizer.json",
22    "tokenizer/tokenizer_config.json",
23    "model.safetensors",
24];
25
26/// Where a name points on the hub.
27#[derive(Debug, Clone, PartialEq, Eq)]
28pub struct HubRef {
29    /// `org/repo`.
30    pub repo: String,
31    /// The folder inside the repo, empty for the root.
32    pub subfolder: String,
33}
34
35impl HubRef {
36    /// Reads an alias or an `hf://org/repo[/subfolder]` reference.
37    #[must_use]
38    pub fn parse(name: &str) -> Option<HubRef> {
39        if let Some((_, repo, sub)) = ALIASES.iter().find(|a| a.0 == name) {
40            return Some(HubRef { repo: (*repo).into(), subfolder: (*sub).into() });
41        }
42        let rest = name.strip_prefix("hf://")?;
43        let mut parts = rest.splitn(3, '/');
44        let (org, repo) = (parts.next()?, parts.next()?);
45        if org.is_empty() || repo.is_empty() {
46            return None;
47        }
48        let subfolder = parts.next().unwrap_or("").trim_matches('/').to_string();
49        Some(HubRef { repo: format!("{org}/{repo}"), subfolder })
50    }
51
52    /// The repo's folder in the cache.
53    #[must_use]
54    pub fn repo_dir(&self, cache: &Path) -> PathBuf {
55        cache.join(format!("models--{}", self.repo.replace('/', "--")))
56    }
57
58    /// The checkpoint folder in the snapshot `refs/main` points at, or in the newest snapshot
59    /// that has it. A repo with several checkpoints in subfolders can have each one pulled at a
60    /// different commit, and `refs/main` only names the last.
61    #[must_use]
62    pub fn local(&self, cache: &Path) -> Option<PathBuf> {
63        let repo = self.repo_dir(cache);
64        let snapshots = repo.join("snapshots");
65        let has = |snap: &Path| {
66            let dir = snap.join(&self.subfolder);
67            dir.join("model.safetensors").is_file().then_some(dir)
68        };
69        if let Ok(rev) = std::fs::read_to_string(repo.join("refs/main"))
70            && let Some(dir) = has(&snapshots.join(rev.trim()))
71        {
72            return Some(dir);
73        }
74        std::fs::read_dir(&snapshots)
75            .ok()?
76            .filter_map(|e| {
77                let e = e.ok()?;
78                let dir = has(&e.path())?;
79                Some((e.metadata().and_then(|m| m.modified()).ok()?, dir))
80            })
81            .max_by_key(|(t, _)| *t)
82            .map(|(_, dir)| dir)
83    }
84}
85
86/// The hub cache: `$HF_HUB_CACHE`, else `$HF_HOME/hub`, else `~/.cache/huggingface/hub`.
87#[must_use]
88pub fn cache_dir() -> PathBuf {
89    let var = |k: &str| std::env::var_os(k).filter(|v| !v.is_empty()).map(PathBuf::from);
90    if let Some(d) = var("HF_HUB_CACHE") {
91        return d;
92    }
93    if let Some(d) = var("HF_HOME") {
94        return d.join("hub");
95    }
96    let home = var("HOME").or_else(|| var("USERPROFILE")).unwrap_or_else(|| PathBuf::from("."));
97    home.join(".cache/huggingface/hub")
98}
99
100/// The path a model name refers to on this machine.
101///
102/// # Errors
103///
104/// A message saying what was looked for and how to fetch it.
105pub fn resolve(name: &str) -> Result<PathBuf, String> {
106    let path = Path::new(name);
107    if path.exists() {
108        return Ok(path.to_path_buf());
109    }
110    let Some(r) = HubRef::parse(name) else {
111        let known: Vec<&str> = ALIASES.iter().map(|a| a.0).collect();
112        return Err(format!(
113            "no model {name:?}: not a path, an hf:// reference or one of {}",
114            known.join(", ")
115        ));
116    };
117    let cache = cache_dir();
118    r.local(&cache)
119        .ok_or_else(|| format!("{name} is not in {}, run kime pull {name} first", cache.display()))
120}
121
122#[cfg(test)]
123mod tests {
124    use super::*;
125
126    #[test]
127    fn names() {
128        let r = HubRef::parse("laya-multilingual").unwrap();
129        assert_eq!(
130            (r.repo.as_str(), r.subfolder.as_str()),
131            ("convaiinnovations/laya", "multilingual")
132        );
133        let r = HubRef::parse("hf://org/repo/a/b/").unwrap();
134        assert_eq!((r.repo.as_str(), r.subfolder.as_str()), ("org/repo", "a/b"));
135        assert_eq!(HubRef::parse("hf://org"), None);
136        assert_eq!(HubRef::parse("kime-v1-s-en"), None);
137        assert_eq!(r.repo_dir(Path::new("/c")), Path::new("/c/models--org--repo"));
138    }
139
140    #[test]
141    fn subfolders_pulled_at_different_commits() {
142        let cache = std::env::temp_dir().join(format!("kime-hub-{}", std::process::id()));
143        let repo = cache.join("models--convaiinnovations--laya");
144        for (rev, sub) in [("old", ""), ("old", "multilingual"), ("new", "typed-decisions")] {
145            let dir = repo.join("snapshots").join(rev).join(sub);
146            std::fs::create_dir_all(&dir).unwrap();
147            std::fs::write(dir.join("model.safetensors"), b"").unwrap();
148        }
149        std::fs::create_dir_all(repo.join("refs")).unwrap();
150        std::fs::write(repo.join("refs/main"), "new").unwrap();
151        let at = |name: &str| HubRef::parse(name).unwrap().local(&cache);
152        assert_eq!(at("laya-typed-decisions"), Some(repo.join("snapshots/new/typed-decisions")));
153        assert_eq!(at("laya-multilingual"), Some(repo.join("snapshots/old/multilingual")));
154        assert_eq!(at("laya"), Some(repo.join("snapshots/old/")));
155        let _ = std::fs::remove_dir_all(cache);
156    }
157}
kime_engine/hub.rs

kime_engine/
hub.rs
