pagesift_ocr/formula_enc/
mod.rs

1//! PP-FormulaNet_plus-L 公式识别的图像编码器（到解码器交叉注意力的 memory 为止）。结构照
2//! transformers 的 `modeling_pp_formulanet.py` 与 `image_processing_pp_formulanet.py`
3//! （Apache-2.0，PaddlePaddle 与 HuggingFace）里的 `PPFormulaNetVisionModel`：
4//!
5//! ```text
6//! 图 ──preprocess──▶ [3, 768, 768]
7//!    ──patch_embed（16×16 卷积，步长 16）──▶ [48·48, 768] + pos_embed
8//!    ──12 层 ViT（窗口 14，第 2、5、8、11 层全局；分解相对位置）──▶ [2304, 768]
9//!    ──neck（1×1 卷积、LN、3×3 卷积、LN）──▶ [256, 48, 48]
10//!    ──multi_modal_projector（两次 3×3 步长 2 卷积、两层 Linear）──▶ [144, 512]
11//! ```
12//!
13//! [`Encoder::encode`] 的输出就是参考实现的 `pooler_output`：形状 `[144, 512]`，
14//! 第 i 行对应 12×12 网格上 (i / 12, i % 12) 处（行优先）、原图约 64×64 像素的一块，
15//! 512 维即解码器的 `d_model`。`PPFormulaNetModel.forward` 把它原样作为 `encoder_hidden_states`
16//! 交给解码器，各层交叉注意力的 k/v 投影直接作用在它上面，中间没有别的变换。
17
18mod attention;
19mod neck;
20mod par;
21mod preprocess;
22mod vit;
23
24use std::path::Path;
25
26use pagesift_infer::{Conv2d, Tensor, Weights, ops};
27
28pub use preprocess::{SIZE, preprocess};
29
30/// [协议] 权重名前缀：`PPFormulaNetForConditionalGeneration.model.encoder`
31const PREFIX: &str = "model.encoder";
32
33/// config.json 里 `vision_config` 用到的部分。
34#[derive(Debug, Clone, PartialEq)]
35pub struct Config {
36    pub image_size: usize,
37    pub patch_size: usize,
38    pub hidden_size: usize,
39    pub layers: usize,
40    pub heads: usize,
41    pub window_size: usize,
42    pub global_attn_indexes: Vec<usize>,
43    pub layer_norm_eps: f32,
44}
45
46impl Config {
47    /// 从 config.json 的 `vision_config` 读。只支持 `use_abs_pos`、`use_rel_pos` 都为真、
48    /// `hidden_act = "gelu"`（PP-FormulaNet_plus-L 就是这样）。
49    pub fn from_json(v: &serde_json::Value) -> Result<Config, String> {
50        let c = &v["vision_config"];
51        let int = |k: &str| {
52            c[k].as_u64()
53                .map(|x| x as usize)
54                .ok_or_else(|| format!("vision_config has no integer {k}"))
55        };
56        for k in ["use_abs_pos", "use_rel_pos"] {
57            if c[k].as_bool() != Some(true) {
58                return Err(format!("only {k} = true is supported"));
59            }
60        }
61        if c["hidden_act"].as_str().unwrap_or("gelu") != "gelu" {
62            return Err("only hidden_act = gelu is supported".into());
63        }
64        let global_attn_indexes = c["global_attn_indexes"]
65            .as_array()
66            .ok_or("vision_config has no global_attn_indexes")?
67            .iter()
68            .map(|x| {
69                x.as_u64()
70                    .map(|x| x as usize)
71                    .ok_or("global_attn_indexes must be integers")
72            })
73            .collect::<Result<_, _>>()?;
74        let cfg = Config {
75            image_size: int("image_size")?,
76            patch_size: int("patch_size")?,
77            hidden_size: int("hidden_size")?,
78            layers: int("num_hidden_layers")?,
79            heads: int("num_attention_heads")?,
80            window_size: int("window_size")?,
81            global_attn_indexes,
82            layer_norm_eps: c["layer_norm_eps"]
83                .as_f64()
84                .ok_or("vision_config has no layer_norm_eps")? as f32,
85        };
86        if cfg.image_size != SIZE {
87            return Err(format!(
88                "preprocessing always outputs {SIZE}, the config says {}",
89                cfg.image_size
90            ));
91        }
92        if cfg.patch_size == 0 || !cfg.image_size.is_multiple_of(cfg.patch_size) {
93            return Err("image_size is not a multiple of patch_size".into());
94        }
95        if cfg.heads == 0 || !cfg.hidden_size.is_multiple_of(cfg.heads) {
96            return Err("hidden_size is not divisible by the number of heads".into());
97        }
98        Ok(cfg)
99    }
100
101    /// 特征图边长（768 / 16 = 48）。
102    pub fn grid(&self) -> usize {
103        self.image_size / self.patch_size
104    }
105}
106
107/// 载入好的编码器。
108pub struct Encoder {
109    cfg: Config,
110    patch: Conv2d,
111    /// `[grid², hidden]`，与 patch 序列逐元素相加
112    pos: Vec<f32>,
113    layers: Vec<vit::Layer>,
114    neck: neck::Neck,
115    proj: neck::Projector,
116}
117
118impl Encoder {
119    /// `dir` 是模型目录：model.safetensors、config.json。
120    pub fn load(dir: &Path) -> Result<Encoder, String> {
121        let cfg = Config::from_json(&crate::read_json(dir, "config.json")?)?;
122        Encoder::from_weights(&crate::read_weights(dir)?, cfg)
123    }
124
125    /// 从已读入的权重建（解码器可共用同一份 [`Weights`]）。用到的参数都复制出来，之后可丢掉 `w`。
126    pub fn from_weights(w: &Weights, cfg: Config) -> Result<Encoder, String> {
127        let (g, c) = (cfg.grid(), cfg.hidden_size);
128        let ps = cfg.patch_size;
129        let patch = Conv2d::load(
130            w,
131            &format!("{PREFIX}.patch_embed.projection"),
132            (ps, ps),
133            (0, 0),
134            (1, 1),
135            1,
136        )?;
137        let pos = w.get(&format!("{PREFIX}.pos_embed"))?;
138        // [协议] pos_embed 是 [1, grid, grid, hidden]（通道在最后，与序列布局一致）
139        if pos.shape != [1, g, g, c] {
140            return Err(format!(
141                "pos_embed has shape {:?}, expected [1, {g}, {g}, {c}]",
142                pos.shape
143            ));
144        }
145        let mut layers = Vec::with_capacity(cfg.layers);
146        for i in 0..cfg.layers {
147            let window = if cfg.global_attn_indexes.contains(&i) {
148                0
149            } else {
150                cfg.window_size
151            };
152            layers.push(vit::Layer::load(
153                w,
154                &format!("{PREFIX}.layers.{i}"),
155                cfg.heads,
156                window,
157                g,
158                cfg.layer_norm_eps,
159            )?);
160        }
161        Ok(Encoder {
162            patch,
163            pos: pos.data.clone(),
164            layers,
165            neck: neck::Neck::load(w, &format!("{PREFIX}.neck"))?,
166            proj: neck::Projector::load(w, &format!("{PREFIX}.multi_modal_projector"))?,
167            cfg,
168        })
169    }
170
171    pub fn config(&self) -> &Config {
172        &self.cfg
173    }
174
175    /// `x` 是 [`preprocess`] 的输出 `[3, 768, 768]`，返回解码器交叉注意力的 memory `[144, 512]`。
176    pub fn encode(&self, x: &Tensor) -> Result<Tensor, String> {
177        self.encode_traced(x, |_, _| {})
178    }
179
180    /// 同 [`encode`](Self::encode)，每一段的输出都交给 `trace(名字, 张量)` 看一眼（检查用）。
181    pub fn encode_traced(
182        &self,
183        x: &Tensor,
184        mut trace: impl FnMut(&str, &Tensor),
185    ) -> Result<Tensor, String> {
186        let s = self.cfg.image_size;
187        // [协议] 输入三通道（`num_channels`）
188        if x.shape != [3, s, s] {
189            return Err(format!("input must be [3, {s}, {s}], got {:?}", x.shape));
190        }
191        let g = self.cfg.grid();
192        let mut h = ops::chw_to_seq(&self.patch.forward(x));
193        for (v, p) in h.data.iter_mut().zip(&self.pos) {
194            *v += p;
195        }
196        trace("patch_embed+pos", &h);
197        for (i, l) in self.layers.iter().enumerate() {
198            h = l.forward(h, g);
199            trace(&format!("layer{i}"), &h);
200        }
201        let n = self.neck.forward(&h, g);
202        trace("neck", &n);
203        let m = self.proj.forward(&n);
204        trace("projector", &m);
205        Ok(m)
206    }
207}
pagesift_ocr/formula_enc/mod.rs

pagesift_ocr/formula_enc/
mod.rs
