laya/
hub.rs

1//! Checkpoint resolution against the Hugging Face Hub.
2
3use hf_hub::{split_id, HFClientBuilder};
4use std::path::{Path, PathBuf};
5
6use crate::agent::{Agent, AgentBuilder};
7use crate::router::{ModelName, Router};
8use crate::{Error, Result};
9
10/// Only these files are fetched, matching upstream's `allow_patterns`.
11const CHECKPOINT_PATTERNS: [&str; 5] = [
12    "model.safetensors",
13    "rl_agent_config.json",
14    "encoder/config.json",
15    "tokenizer/*",
16    "mlx_config.json",
17];
18
19/// Download a checkpoint and return the directory holding it.
20///
21/// A path that exists on disk is returned unchanged, matching upstream's
22/// `resolve_model`.
23pub fn resolve_model(
24    model_id_or_path: &str,
25    subfolder: Option<&str>,
26    revision: Option<&str>,
27    token: Option<String>,
28) -> Result<PathBuf> {
29    if let Some(subfolder) = subfolder {
30        let path = Path::new(subfolder);
31        if path.is_absolute() || path.components().any(|c| c.as_os_str() == "..") {
32            return Err(Error::Config(
33                "subfolder must be a relative path inside the model repository".into(),
34            ));
35        }
36    }
37
38    let local = PathBuf::from(expand_home(model_id_or_path));
39    if local.exists() {
40        return Ok(match subfolder {
41            Some(sub) => local.join(sub),
42            None => local,
43        });
44    }
45    if model_id_or_path.starts_with('/')
46        || model_id_or_path.starts_with("./")
47        || model_id_or_path.starts_with("../")
48        || model_id_or_path.starts_with('~')
49    {
50        return Err(Error::Config(format!(
51            "Local model directory does not exist: {model_id_or_path}"
52        )));
53    }
54
55    let mut builder = HFClientBuilder::new();
56    if let Some(token) = token.or_else(|| std::env::var("HF_TOKEN").ok()) {
57        builder = builder.token(token);
58    }
59    let client = builder
60        .build_sync()
61        .map_err(|error| Error::Hub(error.to_string()))?;
62
63    let (owner, name) = split_id(model_id_or_path);
64    let repository = client.model(owner, name);
65
66    // Restrict to the requested subfolder so the bundle repository does not
67    // pull all three checkpoints when only one is wanted.
68    let patterns: Vec<String> = CHECKPOINT_PATTERNS
69        .iter()
70        .map(|pattern| match subfolder {
71            Some(sub) => format!("{}/{pattern}", sub.trim_end_matches('/')),
72            None => (*pattern).to_string(),
73        })
74        .collect();
75
76    let snapshot = repository
77        .snapshot_download()
78        .maybe_revision(revision.map(str::to_string))
79        .allow_patterns(patterns)
80        .send()
81        .map_err(|error| Error::Hub(format!("{model_id_or_path}: {error}")))?;
82
83    Ok(match subfolder {
84        Some(sub) => snapshot.join(sub),
85        None => snapshot,
86    })
87}
88
89fn expand_home(path: &str) -> String {
90    if let Some(rest) = path.strip_prefix("~/") {
91        if let Some(home) = std::env::var_os("HOME") {
92            return Path::new(&home).join(rest).to_string_lossy().into_owned();
93        }
94    }
95    path.to_string()
96}
97
98impl Agent {
99    /// Load a checkpoint by Hub id or local path, with default settings.
100    pub fn from_pretrained(model_id_or_path: &str) -> Result<Self> {
101        Self::from_pretrained_with(model_id_or_path, None, None, AgentBuilder::new())
102    }
103
104    /// Load a checkpoint, choosing a subfolder, revision and options.
105    pub fn from_pretrained_with(
106        model_id_or_path: &str,
107        subfolder: Option<&str>,
108        revision: Option<&str>,
109        options: AgentBuilder,
110    ) -> Result<Self> {
111        let dir = resolve_model(model_id_or_path, subfolder, revision, None)?;
112        options.build(dir)
113    }
114}
115
116impl Router {
117    /// Download and build a checkpoint, then keep it resident.
118    pub fn load(
119        &mut self,
120        model: ModelName,
121        options: AgentBuilder,
122    ) -> Result<std::sync::Arc<Agent>> {
123        if let Some(agent) = self.get(model) {
124            return Ok(agent);
125        }
126        let (repo, subfolder) = model.bundle_location();
127        let dir = resolve_model(repo, subfolder, None, None)?;
128        let agent = std::sync::Arc::new(options.build(dir)?);
129        self.attach(model, std::sync::Arc::clone(&agent));
130        Ok(agent)
131    }
132
133    /// Download and build several checkpoints up front.
134    ///
135    /// A cold load costs seconds while detection costs microseconds, so a
136    /// server that alternates languages should preload rather than rely on the
137    /// LRU, which would otherwise reload on nearly every request.
138    pub fn preload(
139        &mut self,
140        models: &[ModelName],
141        options: impl Fn() -> AgentBuilder,
142    ) -> Result<()> {
143        self.set_max_loaded_at_least(models.len());
144        for model in models {
145            self.load(*model, options())?;
146        }
147        Ok(())
148    }
149}
laya/hub.rs

laya/
hub.rs
