Skip to main content

pagesift_ocr/
models.rs

1//! 模型清单:从哪下载、锁定的上游版本、每个文件的大小和 SHA-256;本地缓存目录与校验。
2//! 权重是 PaddlePaddle 在 Hugging Face 上发布的 PP-OCRv6 safetensors(Apache-2.0),不随本仓库
3//! 分发,打开 OCR 后第一次用时下载(下载由调用方用自己的 HTTP 客户端做,这里只给地址和校验)。
4
5use std::io::Read;
6use std::path::{Path, PathBuf};
7
8use sha2::{Digest, Sha256};
9
10/// 一个模型文件。
11#[derive(Debug, Clone, Copy, PartialEq, Eq)]
12pub struct File {
13    pub name: &'static str,
14    pub size: u64,
15    /// 小写十六进制的 SHA-256。
16    pub sha256: &'static str,
17}
18
19/// 一个模型(一个 Hugging Face 仓库的一个锁定版本)。
20#[derive(Debug, Clone, Copy, PartialEq, Eq)]
21pub struct Model {
22    /// 仓库名,也是本地缓存里的目录名。
23    pub repo: &'static str,
24    /// 锁定的提交号:内容固定,校验值才有意义。
25    pub revision: &'static str,
26    pub files: &'static [File],
27}
28
29/// [协议] Hugging Face 的组织名
30const ORG: &str = "PaddlePaddle";
31
32impl Model {
33    /// 文件的下载地址(`https://huggingface.co/<组织>/<仓库>/resolve/<提交号>/<文件>`)。
34    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    /// 本地目录:`<缓存根>/<仓库名>`。
42    pub fn dir(&self, root: &Path) -> PathBuf {
43        root.join(self.repo)
44    }
45
46    /// 所有文件都在且大小对(不算哈希,载入模型前用)。
47    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    /// 本地目录里的文件是否齐全、大小和 SHA-256 是否都对;返回第一处问题。
55    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
77/// [策略] 校验时每次读的字节数
78const READ_CHUNK: usize = 1 << 20;
79
80/// 边下载边校验一个文件(模型文件上百 MB,不整个放进内存)。
81pub 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, // [协议] 上游文件大小
93        }
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    /// 大小和 SHA-256 都对才算过;返回问题描述。
102    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
119/// 校验下载来的一个文件。
120pub fn check(f: &File, bytes: &[u8]) -> Result<(), String> {
121    let mut h = Hasher::new(f);
122    h.update(bytes);
123    h.finish()
124}
125
126/// 缓存根目录:`$XDG_CACHE_HOME/pagesift/models`,没设时 `~/.cache/pagesift/models`
127/// (XDG Base Directory 规范)。
128pub 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
136// [协议] 以下大小和 SHA-256 取自各仓库锁定提交的文件(2026-11-03 核对:safetensors 的值与
137// Hugging Face LFS 指针一致,JSON 文件按下载内容计算)
138
139/// CPU 用的小模型:检测。
140pub 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, // [协议] 上游文件大小
147            sha256: "89a96a8adc4e9cd0c994098edc76022e496d35844392562b4694c8fbc583f2da",
148        },
149        File {
150            name: "config.json",
151            size: 1_096, // [协议] 上游文件大小
152            sha256: "b42d76ab6325d5234e1afc9987b73dc592b0ec8437c8a29d46ba6209b79252eb",
153        },
154        File {
155            name: "preprocessor_config.json",
156            size: 838, // [协议] 上游文件大小
157            sha256: "74421569a28a78d417db320f9bc039ce0997a1defe2b35f259dfc74299c9f1ed",
158        },
159    ],
160};
161
162/// CPU 用的小模型:识别。
163pub 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, // [协议] 上游文件大小
170            sha256: "f65a332afe5aa663f0b9d5706f4ae8457b5b4058a842d5c1eb22df505c27d642",
171        },
172        File {
173            name: "config.json",
174            size: 1_083, // [协议] 上游文件大小
175            sha256: "8693fd8485e8543e13d0f6dde3891a3f4af9a47fc4e3a391fbed71f8015c899e",
176        },
177        File {
178            name: "preprocessor_config.json",
179            size: 206_177, // [协议] 上游文件大小
180            sha256: "2b24fa36f548893f26a931cf44f1a1a6b2b14ed7e2f8f8e28e6848801a8278db",
181        },
182    ],
183};
184
185/// GPU 用的中模型:检测。
186pub 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, // [协议] 上游文件大小
193            sha256: "bd393266c02e1a680b1b34c301d5d0d81e6290440b7f8ab0f5d5032276b17eb1",
194        },
195        File {
196            name: "config.json",
197            size: 2_134, // [协议] 上游文件大小
198            sha256: "3ea7e760f64255c152ba139e2f5f798303f7189b87c923cf498e2b2e830cadd3",
199        },
200        File {
201            name: "preprocessor_config.json",
202            size: 838, // [协议] 上游文件大小
203            sha256: "74421569a28a78d417db320f9bc039ce0997a1defe2b35f259dfc74299c9f1ed",
204        },
205    ],
206};
207
208/// GPU 用的中模型:识别。
209pub 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, // [协议] 上游文件大小
216            sha256: "5f43c16f2a684b1d2284662178bdb604febd3d6bfdb5ca73828d08d0f7c0c3e9",
217        },
218        File {
219            name: "config.json",
220            size: 1_131, // [协议] 上游文件大小
221            sha256: "aa5f351b4a7df95cc4953f013d53a6d44205f1bb0aee0ec7eb243d1d8fae1c27",
222        },
223        File {
224            name: "preprocessor_config.json",
225            size: 206_177, // [协议] 上游文件大小
226            sha256: "2b24fa36f548893f26a931cf44f1a1a6b2b14ed7e2f8f8e28e6848801a8278db",
227        },
228    ],
229};
230
231/// 公式识别(只有 L 级有 safetensors)。
232pub 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, // [协议] 上游文件大小
239            sha256: "582848cf9b7066bb1cf3166ad75a476da492c3661ffa61bb7439f9fd098bde13",
240        },
241        File {
242            name: "config.json",
243            size: 1_364, // [协议] 上游文件大小
244            sha256: "3d5c513a525ebef73ba4e319e65de793d95aacc1062e5ce2dd851c7f109c060c",
245        },
246        File {
247            name: "generation_config.json",
248            size: 186, // [协议] 上游文件大小
249            sha256: "e0dc49736b0b06f19cb2074201a072cee23707b7abff4a55a24da19b14abd707",
250        },
251        File {
252            name: "processor_config.json",
253            size: 584, // [协议] 上游文件大小
254            sha256: "f1454999fcb5930c2d93e9dcd03f1adf046245e98e96a7f601dabc6ee01c3b78",
255        },
256        File {
257            name: "tokenizer.json",
258            size: 2_140_014, // [协议] 上游文件大小
259            sha256: "2811d82701ec97c192fa256aa2b4516929373870ae660326cc5b1dc879b95ff2",
260        },
261        File {
262            name: "tokenizer_config.json",
263            size: 4_484, // [协议] 上游文件大小
264            sha256: "0bf3cca5ff1a240004932dbab09373b0e8a651f5f85f2ff6f2f9dec3f0e41947",
265        },
266    ],
267};
268
269/// 版面检测(找公式区域)。
270pub 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, // [协议] 上游文件大小
277            sha256: "5ea422c6cc5fe759a47e1357c35639b58173508e025a3131cbe4b6ac59e2b85e",
278        },
279        File {
280            name: "config.json",
281            size: 2_460, // [协议] 上游文件大小
282            sha256: "3cf834b91d23a756b1519bce4db42c09e852f3e35c35092dd5a3e253a50c071a",
283        },
284        File {
285            name: "inference.yml",
286            size: 1_482, // [协议] 上游文件大小
287            sha256: "506fcfac13b3b546ae40d7886b44126420f392adb694e3f8bb6a6286a1f90fdc",
288        },
289        File {
290            name: "preprocessor_config.json",
291            size: 575, // [协议] 上游文件大小
292            sha256: "519fe0187a43a1ca429e3ad8317bab8700f0d5e8fb3a6e3a0a413ffac078ba42",
293        },
294    ],
295};
296
297#[cfg(test)]
298mod tests {
299    use super::*;
300
301    /// 本机缓存里已有的模型要能通过校验(没有缓存时跳过)。
302    #[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}