monocr_onnx/
model_manager.rs

1//! Model Manager
2//!
3//! This module handles downloading and caching the ONNX model used for OCR.
4//! Models are downloaded from HuggingFace and stored in the user's cache directory.
5
6use indicatif::{ProgressBar, ProgressStyle};
7use reqwest::blocking::Client;
8use std::fs;
9use std::io;
10use std::path::{Path, PathBuf};
11
12/// The Hugging Face repository holding the ONNX artifact.
13pub const MODEL_REPO: &str = "janakhpon/monocr";
14
15/// The pinned revision.
16///
17/// `main` is a moving ref and the artifact has already changed under it: the
18/// model served at one point had a 64-pixel input and 225 output classes, v2
19/// had 128 and 316, and the one served now has 160 and 277. A cache that gates on "the file exists"
20/// cannot tell those apart, so the revision is part of the cache path.
21pub const MODEL_REVISION: &str = "d3d9d5e";
22
23/// Filename of the ONNX model within the repository's `onnx/` directory.
24pub const MODEL_FILENAME: &str = "monocr.onnx";
25
26/// Filename of the charset that belongs to the same revision.
27pub const CHARSET_FILENAME: &str = "charset.txt";
28
29/// Manages downloading and caching of OCR models
30///
31/// This struct handles the lifecycle of the ONNX model file, including:
32/// - Determining the cache location (`~/.monocr/models/<revision>/`)
33/// - Downloading the model and its charset from HuggingFace if not present
34/// - Providing the path to the model file for loading
35pub struct ModelManager {
36    /// Directory where this revision's files are cached
37    cache_dir: PathBuf,
38    /// Base URL for downloading, already pinned to a revision
39    base_url: String,
40    /// Filename of the model file
41    model_filename: String,
42}
43
44impl Default for ModelManager {
45    fn default() -> Self {
46        Self::new()
47    }
48}
49
50impl ModelManager {
51    /// Create a new ModelManager with default settings
52    ///
53    /// # Default Values
54    ///
55    /// - Cache directory: `~/.monocr/models/<revision>/`
56    /// - Base URL: `https://huggingface.co/janakhpon/monocr/resolve/<revision>`
57    /// - Model filename: `monocr.onnx`
58    ///
59    /// # Panics
60    ///
61    /// Panics if the user's home directory cannot be determined
62    pub fn new() -> Self {
63        let home = dirs::home_dir().expect("Failed to get home directory");
64        // Revision-scoped: re-pinning MODEL_REVISION is a cache miss rather
65        // than a silent reuse of whatever was downloaded last time.
66        let cache_dir = home.join(".monocr").join("models").join(MODEL_REVISION);
67
68        Self {
69            cache_dir,
70            base_url: format!("https://huggingface.co/{MODEL_REPO}/resolve/{MODEL_REVISION}"),
71            model_filename: MODEL_FILENAME.to_string(),
72        }
73    }
74
75    /// The directory this manager downloads into.
76    pub fn cache_dir(&self) -> &Path {
77        &self.cache_dir
78    }
79
80    /// The pinned download URL for the ONNX model.
81    pub fn model_url(&self) -> String {
82        format!("{}/onnx/{}", self.base_url, self.model_filename)
83    }
84
85    /// The pinned download URL for the charset.
86    ///
87    /// Same revision as the weights — that is the only way to be sure the two
88    /// agree.
89    pub fn charset_url(&self) -> String {
90        format!("{}/onnx/{}", self.base_url, CHARSET_FILENAME)
91    }
92
93    /// Get the path to the ONNX model file, downloading it if the cache for
94    /// this revision is empty.
95    pub fn get_model_path(&self) -> io::Result<PathBuf> {
96        let model_path = self.cache_dir.join(&self.model_filename);
97
98        if !model_path.exists() {
99            println!(
100                "Model {MODEL_REVISION} not found at {:?}. Downloading...",
101                model_path
102            );
103            self.download(&self.model_url(), &model_path)?;
104            println!("Download complete");
105        }
106
107        Ok(model_path)
108    }
109
110    /// Get the charset published alongside the pinned model.
111    ///
112    /// Preferred over the embedded copy because it comes from the same revision
113    /// as the weights.
114    pub fn get_charset(&self) -> io::Result<String> {
115        let charset_path = self.cache_dir.join(CHARSET_FILENAME);
116
117        if !charset_path.exists() {
118            self.download(&self.charset_url(), &charset_path)?;
119        }
120
121        fs::read_to_string(&charset_path)
122    }
123
124    /// Download `url` to `dest` via a temporary file.
125    ///
126    /// The rename is the last step, so an interrupted transfer never leaves a
127    /// truncated artifact behind for the existence check to accept.
128    fn download(&self, url: &str, dest: &Path) -> io::Result<()> {
129        if let Some(parent) = dest.parent() {
130            fs::create_dir_all(parent)?;
131        }
132
133        let client = Client::new();
134        let mut response = client
135            .get(url)
136            .send()
137            .map_err(|e| io::Error::other(format!("Failed to download from {}: {}", url, e)))?;
138
139        if !response.status().is_success() {
140            return Err(io::Error::other(format!(
141                "Failed to download {}: {}",
142                url,
143                response.status()
144            )));
145        }
146
147        let total_size = response.content_length().unwrap_or(0);
148        let pb = ProgressBar::new(total_size);
149        pb.set_style(ProgressStyle::default_bar()
150            .template("{spinner:.green} [{elapsed_precise}] [{wide_bar:.cyan/blue}] {bytes}/{total_bytes} ({eta})")
151            .unwrap()
152            .progress_chars(">-"));
153
154        let tmp_path = dest.with_extension("part");
155        let mut file = fs::File::create(&tmp_path)?;
156        let copied = io::copy(&mut response, &mut file);
157        drop(file);
158
159        if let Err(e) = copied {
160            let _ = fs::remove_file(&tmp_path);
161            return Err(e);
162        }
163
164        fs::rename(&tmp_path, dest)?;
165        pb.finish_and_clear();
166        Ok(())
167    }
168}
169
170#[cfg(test)]
171mod tests {
172    use super::*;
173
174    /// `main` is a moving ref and the artifact has already changed under it.
175    #[test]
176    fn download_urls_are_pinned() {
177        let m = ModelManager::new();
178        for url in [m.model_url(), m.charset_url()] {
179            assert!(
180                !url.contains("/resolve/main/"),
181                "still tracking the moving ref `main`: {url}"
182            );
183            assert!(
184                url.contains(&format!("/resolve/{MODEL_REVISION}/")),
185                "not pinned to {MODEL_REVISION}: {url}"
186            );
187        }
188    }
189
190    /// The charset has to come from the same revision as the weights.
191    #[test]
192    fn charset_is_fetched_from_the_model_revision() {
193        let m = ModelManager::new();
194        let model_dir = m.model_url().trim_end_matches(MODEL_FILENAME).to_string();
195        let charset_dir = m
196            .charset_url()
197            .trim_end_matches(CHARSET_FILENAME)
198            .to_string();
199        assert_eq!(model_dir, charset_dir);
200    }
201
202    /// The cache used to gate on file existence alone, so an artifact from an
203    /// older revision was reused forever.
204    #[test]
205    fn cache_dir_is_scoped_by_revision() {
206        let m = ModelManager::new();
207        assert_eq!(
208            m.cache_dir().file_name().and_then(|s| s.to_str()),
209            Some(MODEL_REVISION),
210            "cache directory {:?} is not revision-scoped",
211            m.cache_dir()
212        );
213    }
214}
monocr_onnx/model_manager.rs

monocr_onnx/
model_manager.rs
