ffai_models/
lib.rs

1//! # ffai-models
2//!
3//! Weight management for `FFai`. Core principle: **weights are data, not
4//! code** — the repo holds only TOML *manifests* describing each model (files,
5//! source, checksums, and crucially its *license*, which is often different
6//! from `FFai`'s). Weights are fetched into a local cache, never vendored.
7//!
8//! Phase 0 ships manifests + cache resolution; the downloader (Hugging Face
9//! hub, resumable, checksum-verified) lands in Phase 1.
10
11use std::collections::BTreeMap;
12use std::path::{Path, PathBuf};
13
14use ffai_core::error::{Error, Result};
15use serde::Deserialize;
16use sha2::{Digest, Sha256};
17
18/// A model whose files are all present locally, with their resolved paths.
19#[derive(Debug, Clone)]
20pub struct ResolvedModel {
21    pub name: String,
22    /// The WEIGHTS' license — may be more restrictive than `FFai`'s own.
23    pub license: String,
24    /// Filename → local path.
25    pub files: BTreeMap<String, PathBuf>,
26}
27
28impl ResolvedModel {
29    /// Path of a required file, or a clear error naming what's missing.
30    pub fn file(&self, name: &str) -> Result<&Path> {
31        self.files
32            .get(name)
33            .map(PathBuf::as_path)
34            .ok_or_else(|| Error::Model(format!("model `{}` has no file `{name}`", self.name)))
35    }
36}
37
38/// Resolve one file from the shared Hugging Face cache, downloading it when
39/// `cache_only` is false. Uses the same cache as `transformers` and
40/// `faster-whisper`, so a model either side already has is not fetched twice.
41#[cfg(feature = "fetch")]
42fn hub_download(repo: &str, filename: &str, cache_only: bool) -> Result<PathBuf> {
43    let (owner, name) = repo
44        .split_once('/')
45        .ok_or_else(|| Error::Model(format!("hf_repo `{repo}` is not in `owner/name` form")))?;
46    let client = hf_hub::HFClientSync::new()
47        .map_err(|e| Error::Model(format!("hugging face client init failed: {e}")))?;
48    client
49        .model(owner, name)
50        .download_file()
51        .filename(filename)
52        .local_files_only(cache_only)
53        .send()
54        .map_err(|e| Error::Model(format!("{repo}/{filename}: {e}")))
55}
56
57/// Without `fetch`, nothing is downloaded — and the error says which file to
58/// supply rather than failing as a missing symbol.
59#[cfg(not(feature = "fetch"))]
60fn hub_download(repo: &str, filename: &str, _cache_only: bool) -> Result<PathBuf> {
61    Err(Error::Model(format!(
62        "{repo}/{filename} is not present locally and this build has the          `fetch` feature disabled — place the file in the model directory, or          enable `ffai-models/fetch` to download it"
63    )))
64}
65
66/// Verify a downloaded file against its manifest checksum, when one is
67/// declared. A mismatch is an error, never a warning: silently running on
68/// unexpected weights would invalidate every measurement taken with them.
69fn verify_checksum(path: &Path, file: &ModelFile) -> Result<()> {
70    let Some(expected) = &file.sha256 else {
71        return Ok(());
72    };
73    let bytes = std::fs::read(path)?;
74    let actual: String = Sha256::digest(&bytes)
75        .iter()
76        .map(|b| format!("{b:02x}"))
77        .collect();
78    if actual != expected.to_ascii_lowercase() {
79        return Err(Error::Model(format!(
80            "checksum mismatch for {}: manifest says {expected}, file is {actual}",
81            path.display()
82        )));
83    }
84    Ok(())
85}
86
87/// One weight/config file belonging to a model.
88#[derive(Debug, Clone, Deserialize)]
89pub struct ModelFile {
90    /// Filename within the model's cache directory (and its HF repo).
91    pub name: String,
92    /// Hex SHA-256, verified after download when present.
93    pub sha256: Option<String>,
94}
95
96/// A model manifest (`models/*.toml`).
97#[derive(Debug, Clone, Deserialize)]
98pub struct ModelManifest {
99    pub name: String,
100    /// Task tag: "asr", "tts", "ocr", "vlm".
101    pub task: String,
102    pub description: Option<String>,
103    /// The WEIGHTS' license — surfaced to users because it may be more
104    /// restrictive than `FFai`'s MIT/Apache code license.
105    pub license: String,
106    /// Hugging Face repo id, e.g. "openai/whisper-tiny".
107    pub hf_repo: Option<String>,
108    #[serde(default)]
109    pub files: Vec<ModelFile>,
110}
111
112impl ModelManifest {
113    pub fn from_toml(text: &str) -> Result<Self> {
114        toml::from_str(text).map_err(|e| Error::Model(format!("bad manifest: {e}")))
115    }
116
117    pub fn load(path: &Path) -> Result<Self> {
118        let text = std::fs::read_to_string(path)?;
119        Self::from_toml(&text)
120    }
121
122    /// Directory this model's files live in when placed manually.
123    #[must_use]
124    pub fn cache_path(&self) -> PathBuf {
125        cache_dir().join("models").join(&self.name)
126    }
127
128    /// True when every listed file already resolves locally — no network.
129    #[must_use]
130    pub fn is_cached(&self) -> bool {
131        !self.files.is_empty()
132            && self
133                .files
134                .iter()
135                .all(|f| self.local_path(&f.name).is_some())
136    }
137
138    /// Resolve one file without touching the network: a manual placement
139    /// under [`Self::cache_path`] wins, otherwise the shared Hugging Face
140    /// cache (the same one `transformers`/`faster-whisper` use, so a model
141    /// downloaded by either side is not downloaded twice).
142    #[must_use]
143    pub fn local_path(&self, name: &str) -> Option<PathBuf> {
144        let manual = self.cache_path().join(name);
145        if manual.exists() {
146            return Some(manual);
147        }
148        hub_download(self.hf_repo.as_ref()?, name, true).ok()
149    }
150
151    /// Resolve every file, downloading from the Hugging Face hub as needed.
152    ///
153    /// Downloads are cached, so this is cheap on repeat calls — but it is
154    /// still network I/O the first time, which is why benchmarks warm the
155    /// cache outside any timed region (see docs/benchmarking.md).
156    pub fn fetch(&self) -> Result<ResolvedModel> {
157        let mut files = BTreeMap::new();
158        for file in &self.files {
159            if let Some(path) = self.local_path(&file.name) {
160                verify_checksum(&path, file)?;
161                files.insert(file.name.clone(), path);
162                continue;
163            }
164            let repo = self.hf_repo.as_ref().ok_or_else(|| {
165                Error::Model(format!(
166                    "model `{}` declares no hf_repo and `{}` is not present under {}",
167                    self.name,
168                    file.name,
169                    self.cache_path().display()
170                ))
171            })?;
172            let path = hub_download(repo, &file.name, false)?;
173            verify_checksum(&path, file)?;
174            files.insert(file.name.clone(), path);
175        }
176        Ok(ResolvedModel {
177            name: self.name.clone(),
178            license: self.license.clone(),
179            files,
180        })
181    }
182}
183
184/// Load every `*.toml` manifest in a directory (typically `models/`).
185pub fn load_dir(dir: &Path) -> Result<Vec<ModelManifest>> {
186    let mut out = Vec::new();
187    for entry in std::fs::read_dir(dir)? {
188        let path = entry?.path();
189        if path.extension().and_then(|e| e.to_str()) == Some("toml") {
190            out.push(ModelManifest::load(&path)?);
191        }
192    }
193    out.sort_by(|a, b| a.task.cmp(&b.task).then_with(|| a.name.cmp(&b.name)));
194    Ok(out)
195}
196
197/// The `FFai` model cache root: `$FFAI_CACHE` or `<os cache dir>/ffai`.
198pub fn cache_dir() -> PathBuf {
199    if let Ok(dir) = std::env::var("FFAI_CACHE") {
200        return PathBuf::from(dir);
201    }
202    dirs::cache_dir()
203        .unwrap_or_else(std::env::temp_dir)
204        .join("ffai")
205}
206
207#[cfg(test)]
208mod tests {
209    use super::*;
210
211    #[test]
212    fn manifest_parses_and_surfaces_license() {
213        let m = ModelManifest::from_toml(
214            r#"
215            name = "whisper-tiny"
216            task = "asr"
217            license = "Apache-2.0"
218            hf_repo = "openai/whisper-tiny"
219
220            [[files]]
221            name = "model.safetensors"
222            "#,
223        )
224        .unwrap();
225        assert_eq!(m.name, "whisper-tiny");
226        assert_eq!(m.license, "Apache-2.0");
227        assert_eq!(m.files.len(), 1);
228        assert!(!m.is_cached());
229    }
230}
ffai_models/lib.rs

ffai_models/
lib.rs
