1use std::io::Read;
6use std::path::{Path, PathBuf};
7
8use sha2::{Digest, Sha256};
9
10#[derive(Debug, Clone, Copy, PartialEq, Eq)]
12pub struct File {
13 pub name: &'static str,
14 pub size: u64,
15 pub sha256: &'static str,
17}
18
19#[derive(Debug, Clone, Copy, PartialEq, Eq)]
21pub struct Model {
22 pub repo: &'static str,
24 pub revision: &'static str,
26 pub files: &'static [File],
27}
28
29const ORG: &str = "PaddlePaddle";
31
32impl Model {
33 pub fn url(&self, file: &File) -> String {
35 format!(
36 "https://huggingface.co/{ORG}/{}/resolve/{}/{}",
37 self.repo, self.revision, file.name
38 )
39 }
40
41 pub fn dir(&self, root: &Path) -> PathBuf {
43 root.join(self.repo)
44 }
45
46 pub fn present(&self, root: &Path) -> bool {
48 let dir = self.dir(root);
49 self.files.iter().all(|f| {
50 std::fs::metadata(dir.join(f.name)).is_ok_and(|m| m.is_file() && m.len() == f.size)
51 })
52 }
53
54 pub fn verify(&self, root: &Path) -> Result<(), String> {
56 let dir = self.dir(root);
57 for f in self.files {
58 let p = dir.join(f.name);
59 let mut hash = Hasher::new(f);
60 let mut file = std::fs::File::open(&p).map_err(|e| format!("{}: {e}", p.display()))?;
61 let mut buf = vec![0u8; READ_CHUNK];
62 loop {
63 let n = file
64 .read(&mut buf)
65 .map_err(|e| format!("{}: {e}", p.display()))?;
66 if n == 0 {
67 break;
68 }
69 hash.update(&buf[..n]);
70 }
71 hash.finish().map_err(|e| format!("{}: {e}", p.display()))?;
72 }
73 Ok(())
74 }
75}
76
77const READ_CHUNK: usize = 1 << 20;
79
80pub struct Hasher<'a> {
82 file: &'a File,
83 sha: Sha256,
84 size: u64,
85}
86
87impl<'a> Hasher<'a> {
88 pub fn new(file: &'a File) -> Self {
89 Self {
90 file,
91 sha: Sha256::new(),
92 size: 0, }
94 }
95
96 pub fn update(&mut self, bytes: &[u8]) {
97 self.sha.update(bytes);
98 self.size += bytes.len() as u64;
99 }
100
101 pub fn finish(self) -> Result<(), String> {
103 if self.size != self.file.size {
104 return Err(format!("{} bytes, expected {}", self.size, self.file.size));
105 }
106 let got: String = self
107 .sha
108 .finalize()
109 .iter()
110 .map(|b| format!("{b:02x}"))
111 .collect();
112 if got != self.file.sha256 {
113 return Err(format!("SHA-256 is {got}, expected {}", self.file.sha256));
114 }
115 Ok(())
116 }
117}
118
119pub fn check(f: &File, bytes: &[u8]) -> Result<(), String> {
121 let mut h = Hasher::new(f);
122 h.update(bytes);
123 h.finish()
124}
125
126pub fn cache_root() -> Option<PathBuf> {
129 let base = std::env::var_os("XDG_CACHE_HOME")
130 .filter(|v| !v.is_empty())
131 .map(PathBuf::from)
132 .or_else(|| std::env::var_os("HOME").map(|h| PathBuf::from(h).join(".cache")))?;
133 Some(base.join("pagesift").join("models"))
134}
135
136pub const SMALL_DET: Model = Model {
141 repo: "PP-OCRv6_small_det_safetensors",
142 revision: "eae2ee920a39fb3087637d3dbb58df1896ec1f24",
143 files: &[
144 File {
145 name: "model.safetensors",
146 size: 9_938_124, sha256: "89a96a8adc4e9cd0c994098edc76022e496d35844392562b4694c8fbc583f2da",
148 },
149 File {
150 name: "config.json",
151 size: 1_096, sha256: "b42d76ab6325d5234e1afc9987b73dc592b0ec8437c8a29d46ba6209b79252eb",
153 },
154 File {
155 name: "preprocessor_config.json",
156 size: 838, sha256: "74421569a28a78d417db320f9bc039ce0997a1defe2b35f259dfc74299c9f1ed",
158 },
159 ],
160};
161
162pub const SMALL_REC: Model = Model {
164 repo: "PP-OCRv6_small_rec_safetensors",
165 revision: "fe049fb103f57443fe8840c54ed06b702f3c1de5",
166 files: &[
167 File {
168 name: "model.safetensors",
169 size: 21_204_736, sha256: "f65a332afe5aa663f0b9d5706f4ae8457b5b4058a842d5c1eb22df505c27d642",
171 },
172 File {
173 name: "config.json",
174 size: 1_083, sha256: "8693fd8485e8543e13d0f6dde3891a3f4af9a47fc4e3a391fbed71f8015c899e",
176 },
177 File {
178 name: "preprocessor_config.json",
179 size: 206_177, sha256: "2b24fa36f548893f26a931cf44f1a1a6b2b14ed7e2f8f8e28e6848801a8278db",
181 },
182 ],
183};
184
185pub const MEDIUM_DET: Model = Model {
187 repo: "PP-OCRv6_medium_det_safetensors",
188 revision: "4236c2b61741a259c091fd879dcc4edc339e916c",
189 files: &[
190 File {
191 name: "model.safetensors",
192 size: 88_020_412, sha256: "bd393266c02e1a680b1b34c301d5d0d81e6290440b7f8ab0f5d5032276b17eb1",
194 },
195 File {
196 name: "config.json",
197 size: 2_134, sha256: "3ea7e760f64255c152ba139e2f5f798303f7189b87c923cf498e2b2e830cadd3",
199 },
200 File {
201 name: "preprocessor_config.json",
202 size: 838, sha256: "74421569a28a78d417db320f9bc039ce0997a1defe2b35f259dfc74299c9f1ed",
204 },
205 ],
206};
207
208pub const MEDIUM_REC: Model = Model {
210 repo: "PP-OCRv6_medium_rec_safetensors",
211 revision: "024cad6a831de75c2c3c26e711ba8c4a82ccd24b",
212 files: &[
213 File {
214 name: "model.safetensors",
215 size: 76_741_720, sha256: "5f43c16f2a684b1d2284662178bdb604febd3d6bfdb5ca73828d08d0f7c0c3e9",
217 },
218 File {
219 name: "config.json",
220 size: 1_131, sha256: "aa5f351b4a7df95cc4953f013d53a6d44205f1bb0aee0ec7eb243d1d8fae1c27",
222 },
223 File {
224 name: "preprocessor_config.json",
225 size: 206_177, sha256: "2b24fa36f548893f26a931cf44f1a1a6b2b14ed7e2f8f8e28e6848801a8278db",
227 },
228 ],
229};
230
231pub const FORMULA: Model = Model {
233 repo: "PP-FormulaNet_plus-L_safetensors",
234 revision: "6cc004a8e91ea3f3819b05317f40b290cbe99c2b",
235 files: &[
236 File {
237 name: "model.safetensors",
238 size: 727_633_048, sha256: "582848cf9b7066bb1cf3166ad75a476da492c3661ffa61bb7439f9fd098bde13",
240 },
241 File {
242 name: "config.json",
243 size: 1_364, sha256: "3d5c513a525ebef73ba4e319e65de793d95aacc1062e5ce2dd851c7f109c060c",
245 },
246 File {
247 name: "generation_config.json",
248 size: 186, sha256: "e0dc49736b0b06f19cb2074201a072cee23707b7abff4a55a24da19b14abd707",
250 },
251 File {
252 name: "processor_config.json",
253 size: 584, sha256: "f1454999fcb5930c2d93e9dcd03f1adf046245e98e96a7f601dabc6ee01c3b78",
255 },
256 File {
257 name: "tokenizer.json",
258 size: 2_140_014, sha256: "2811d82701ec97c192fa256aa2b4516929373870ae660326cc5b1dc879b95ff2",
260 },
261 File {
262 name: "tokenizer_config.json",
263 size: 4_484, sha256: "0bf3cca5ff1a240004932dbab09373b0e8a651f5f85f2ff6f2f9dec3f0e41947",
265 },
266 ],
267};
268
269pub const LAYOUT: Model = Model {
271 repo: "PP-DocLayoutV3_safetensors",
272 revision: "e4f489dc536556fc1be02b973ba06756e0d4a2ac",
273 files: &[
274 File {
275 name: "model.safetensors",
276 size: 133_270_468, sha256: "5ea422c6cc5fe759a47e1357c35639b58173508e025a3131cbe4b6ac59e2b85e",
278 },
279 File {
280 name: "config.json",
281 size: 2_460, sha256: "3cf834b91d23a756b1519bce4db42c09e852f3e35c35092dd5a3e253a50c071a",
283 },
284 File {
285 name: "inference.yml",
286 size: 1_482, sha256: "506fcfac13b3b546ae40d7886b44126420f392adb694e3f8bb6a6286a1f90fdc",
288 },
289 File {
290 name: "preprocessor_config.json",
291 size: 575, sha256: "519fe0187a43a1ca429e3ad8317bab8700f0d5e8fb3a6e3a0a413ffac078ba42",
293 },
294 ],
295};
296
297#[cfg(test)]
298mod tests {
299 use super::*;
300
301 #[test]
303 fn cached_models_verify() {
304 let Some(root) = cache_root() else { return };
305 for m in [
306 SMALL_DET, SMALL_REC, MEDIUM_DET, MEDIUM_REC, FORMULA, LAYOUT,
307 ] {
308 if m.dir(&root).join("model.safetensors").exists() {
309 m.verify(&root).unwrap();
310 }
311 }
312 }
313
314 #[test]
315 fn rejects_wrong_bytes() {
316 let f = &SMALL_DET.files[1];
317 assert!(check(f, b"{}").is_err());
318 }
319}