pagesift_ocr/formula/
decoder.rs

1//! mBART 式文本解码器与贪心生成，对应参考实现的 `PPFormulaNetTextModel`、
2//! `PPFormulaNetDecoderLayer` 和 `PPFormulaNetForConditionalGeneration`：
3//!
4//! 词向量 × √d_model + 位置向量（下标偏移 2）→ layernorm_embedding →
5//! 每层（前置 LayerNorm）：因果自注意力、对 memory 的交叉注意力、前馈，都带残差 →
6//! 末尾 layer_norm → lm_head（无偏置）。
7//!
8//! 自注意力的 K、V 逐步追加，交叉注意力的 K、V 每层只算一次（见 [`Cache`]）。
9
10use std::path::Path;
11
12use pagesift_infer::{Activation, LayerNorm, Linear, Tensor, Weights, ops};
13use serde_json::Value;
14
15use super::attention::Attention;
16
17/// [协议] `nn.LayerNorm` 的缺省 eps（参考实现没有改它）
18const LN_EPS: f32 = 1e-5;
19/// [协议] `PPFormulaNetLearnedPositionalEmbedding.offset`：位置下标统一加 2
20const POSITION_OFFSET: usize = 2;
21/// [协议] 权重名前缀：`PPFormulaNetForConditionalGeneration.model.decoder`
22const PREFIX: &str = "model.decoder";
23
24struct Layer {
25    self_ln: LayerNorm,
26    self_attn: Attention,
27    cross_ln: LayerNorm,
28    cross_attn: Attention,
29    final_ln: LayerNorm,
30    fc1: Linear,
31    fc2: Linear,
32    act: Activation,
33}
34
35/// 一层的缓存：自注意力的 K、V（`[已生成, D]`，逐步追加），交叉注意力的 K、V（`[L, D]`，只算一次）。
36struct LayerCache {
37    k: Vec<f32>,
38    v: Vec<f32>,
39    cross_k: Vec<f32>,
40    cross_v: Vec<f32>,
41}
42
43/// 一次生成的缓存。
44pub(super) struct Cache {
45    layers: Vec<LayerCache>,
46    /// 已经喂进去的 token 数，也是下一个 token 的位置。
47    len: usize,
48    /// memory 的长度 L。
49    memory_len: usize,
50}
51
52impl Layer {
53    fn load(w: &Weights, p: &str, heads: usize, act: Activation) -> Result<Layer, String> {
54        Ok(Layer {
55            self_ln: LayerNorm::load(w, &format!("{p}.self_attn_layer_norm"), LN_EPS)?,
56            self_attn: Attention::load(w, &format!("{p}.self_attn"), heads)?,
57            cross_ln: LayerNorm::load(w, &format!("{p}.encoder_attn_layer_norm"), LN_EPS)?,
58            cross_attn: Attention::load(w, &format!("{p}.encoder_attn"), heads)?,
59            final_ln: LayerNorm::load(w, &format!("{p}.final_layer_norm"), LN_EPS)?,
60            fc1: Linear::load(w, &format!("{p}.fc1"))?,
61            fc2: Linear::load(w, &format!("{p}.fc2"))?,
62            act,
63        })
64    }
65
66    /// `x` 是新喂入的 `[T, D]`，`past` 是这之前已缓存的长度。
67    fn forward(&self, x: Tensor, c: &mut LayerCache, past: usize, memory_len: usize) -> Tensor {
68        let n = self.self_ln.forward(&x);
69        c.k.extend_from_slice(&self.self_attn.k.forward(&n).data);
70        c.v.extend_from_slice(&self.self_attn.v.forward(&n).data);
71        let a = self.self_attn.attend(&n, &c.k, &c.v, |i| past + i + 1);
72        let x = ops::add(&x, &a);
73
74        let n = self.cross_ln.forward(&x);
75        let a = self
76            .cross_attn
77            .attend(&n, &c.cross_k, &c.cross_v, |_| memory_len);
78        let x = ops::add(&x, &a);
79
80        let mut h = self.fc1.forward(&self.final_ln.forward(&x));
81        self.act.apply_all(&mut h);
82        ops::add(&x, &self.fc2.forward(&h))
83    }
84}
85
86/// 载入好的解码器。
87pub struct Decoder {
88    /// `[词表, D]`
89    embed_tokens: Tensor,
90    /// `[最大位置 + 2, D]`
91    embed_positions: Tensor,
92    embed_scale: f32,
93    layernorm_embedding: LayerNorm,
94    layers: Vec<Layer>,
95    layer_norm: LayerNorm,
96    lm_head: Linear,
97    d_model: usize,
98    start: u32,
99    eos: u32,
100    forced_eos: Option<u32>,
101    max_length: usize,
102}
103
104fn cfg_usize(c: &Value, key: &str) -> Result<usize, String> {
105    c[key]
106        .as_u64()
107        .map(|v| v as usize)
108        .ok_or_else(|| format!("the config has no {key}"))
109}
110
111impl Decoder {
112    /// `dir` 是模型目录：model.safetensors、config.json、generation_config.json。
113    pub fn load(dir: &Path) -> Result<Decoder, String> {
114        let cfg = crate::read_json(dir, "config.json")?;
115        let generation = crate::read_json(dir, "generation_config.json")?;
116        let w = crate::read_weights(dir)?;
117        Decoder::from_weights(&w, &cfg["text_config"], &generation)
118    }
119
120    /// 从已读入的权重构造（和图像编码器共用一次读文件时用）。
121    /// `text` 是 config.json 的 `text_config`，`generation` 是 generation_config.json。
122    pub fn from_weights(w: &Weights, text: &Value, generation: &Value) -> Result<Decoder, String> {
123        let d_model = cfg_usize(text, "d_model")?;
124        let heads = cfg_usize(text, "decoder_attention_heads")?;
125        if heads == 0 || !d_model.is_multiple_of(heads) {
126            return Err(format!(
127                "d_model {d_model} is not divisible by {heads} heads"
128            ));
129        }
130        let act = match text["activation_function"].as_str().unwrap_or("") {
131            "gelu" => Activation::Gelu,
132            "relu" => Activation::Relu,
133            "silu" | "swish" => Activation::Silu,
134            a => return Err(format!("unsupported activation {a:?}")),
135        };
136        let layers = (0..cfg_usize(text, "decoder_layers")?)
137            .map(|i| Layer::load(w, &format!("{PREFIX}.layers.{i}"), heads, act))
138            .collect::<Result<Vec<_>, _>>()?;
139        let embed_tokens = w.get(&format!("{PREFIX}.embed_tokens.weight"))?.clone();
140        let embed_positions = w.get(&format!("{PREFIX}.embed_positions.weight"))?.clone();
141        if embed_tokens.tc().1 != d_model || embed_positions.tc().1 != d_model {
142            return Err(format!(
143                "token or position embedding size does not match d_model {d_model}"
144            ));
145        }
146        let embed_scale = if text["scale_embedding"].as_bool().unwrap_or(false) {
147            (d_model as f32).sqrt()
148        } else {
149            1.0
150        };
151        let token = |key: &str| cfg_usize(generation, key).map(|v| v as u32);
152        Ok(Decoder {
153            embed_tokens,
154            embed_positions,
155            embed_scale,
156            layernorm_embedding: LayerNorm::load(
157                w,
158                &format!("{PREFIX}.layernorm_embedding"),
159                LN_EPS,
160            )?,
161            layers,
162            layer_norm: LayerNorm::load(w, &format!("{PREFIX}.layer_norm"), LN_EPS)?,
163            lm_head: Linear::load(w, "lm_head")?,
164            d_model,
165            start: token("decoder_start_token_id")?,
166            eos: token("eos_token_id")?,
167            forced_eos: token("forced_eos_token_id").ok(),
168            max_length: cfg_usize(generation, "max_length")?,
169        })
170    }
171
172    /// generation_config.json 的 `max_length`（含起始 token 的序列总长），传给 [`generate`](Self::generate)。
173    pub fn max_length(&self) -> usize {
174        self.max_length
175    }
176
177    /// 贪心生成。`memory` 是编码器输出 `[L, d_model]`（参考实现的 `encoder_hidden_states`）。
178    ///
179    /// `max_len` 的含义同 `generate(max_length=…)`：序列总长的上限，含起始 token；序列长到
180    /// `max_len − 1` 时下一个 token 强制为 `forced_eos_token_id`
181    /// （`ForcedEOSTokenLogitsProcessor`）。超过位置向量表的长度时取表长。
182    ///
183    /// 返回生成的 token，不含起始 token 和结尾的 eos。
184    pub fn generate(&self, memory: &Tensor, max_len: usize) -> Result<Vec<u32>, String> {
185        let max_len = max_len.min(self.max_positions());
186        let mut cache = self.new_cache(memory)?;
187        let mut last = self.start;
188        let mut out = Vec::new();
189        // 序列长度 = 起始 token + 已生成的
190        for len in 1..max_len {
191            let hidden = self.forward(&[last], &mut cache)?;
192            let next = match self.forced_eos {
193                Some(f) if len + 1 == max_len => f,
194                _ => argmax(&self.logits(&hidden)?),
195            };
196            if next == self.eos {
197                break;
198            }
199            out.push(next);
200            last = next;
201        }
202        Ok(out)
203    }
204
205    /// 位置向量表能容纳的序列长度。
206    fn max_positions(&self) -> usize {
207        self.embed_positions.tc().0.saturating_sub(POSITION_OFFSET)
208    }
209
210    /// 建一次生成的缓存：交叉注意力的 K、V 每层在这里算一次。
211    pub(super) fn new_cache(&self, memory: &Tensor) -> Result<Cache, String> {
212        if memory.shape.len() != 2 || memory.shape[1] != self.d_model || memory.shape[0] == 0 {
213            return Err(format!(
214                "memory must be [L, {}], got {:?}",
215                self.d_model, memory.shape
216            ));
217        }
218        let layers = self
219            .layers
220            .iter()
221            .map(|l| LayerCache {
222                k: Vec::new(),
223                v: Vec::new(),
224                cross_k: l.cross_attn.k.forward(memory).data,
225                cross_v: l.cross_attn.v.forward(memory).data,
226            })
227            .collect();
228        Ok(Cache {
229            layers,
230            len: 0,
231            memory_len: memory.shape[0],
232        })
233    }
234
235    /// 喂入新的 token（接在缓存之后），返回这些 token 的隐状态 `[T, D]`。
236    pub(super) fn forward(&self, ids: &[u32], cache: &mut Cache) -> Result<Tensor, String> {
237        let (vocab, d) = self.embed_tokens.tc();
238        let (past, memory_len) = (cache.len, cache.memory_len);
239        if past + ids.len() > self.max_positions() {
240            return Err(format!(
241                "sequence length {} exceeds the position table of {}",
242                past + ids.len(),
243                self.max_positions()
244            ));
245        }
246        let mut x = Vec::with_capacity(ids.len() * d);
247        for (i, &id) in ids.iter().enumerate() {
248            let id = id as usize;
249            if id >= vocab {
250                return Err(format!("token {id} is outside the vocabulary of {vocab}"));
251            }
252            let tok = &self.embed_tokens.data[id * d..(id + 1) * d];
253            let at = (past + i + POSITION_OFFSET) * d;
254            let pos = &self.embed_positions.data[at..at + d];
255            x.extend(tok.iter().zip(pos).map(|(t, p)| t * self.embed_scale + p));
256        }
257        let mut x = self
258            .layernorm_embedding
259            .forward(&Tensor::new(vec![ids.len(), d], x));
260        for (layer, c) in self.layers.iter().zip(&mut cache.layers) {
261            x = layer.forward(x, c, past, memory_len);
262        }
263        cache.len += ids.len();
264        Ok(self.layer_norm.forward(&x))
265    }
266
267    /// 隐状态最后一行的 logits。出现 NaN 或无穷时报错，不让它变成一串看似正常的 token。
268    pub(super) fn logits(&self, hidden: &Tensor) -> Result<Vec<f32>, String> {
269        let (t, d) = hidden.tc();
270        let last = Tensor::new(vec![1, d], hidden.data[(t - 1) * d..].to_vec());
271        let y = self.lm_head.forward(&last).data;
272        if y.iter().any(|v| !v.is_finite()) {
273            return Err("NaN or infinity in the logits".into());
274        }
275        Ok(y)
276    }
277}
278
279/// 最大值的下标，相同时取最前面的（`torch.argmax`）。
280fn argmax(v: &[f32]) -> u32 {
281    let (mut arg, mut m) = (0, f32::NEG_INFINITY);
282    for (i, &x) in v.iter().enumerate() {
283        if x > m {
284            (arg, m) = (i, x);
285        }
286    }
287    arg as u32
288}
289
290#[cfg(test)]
291#[path = "decoder_tests.rs"]
292mod tests;
pagesift_ocr/formula/decoder.rs

pagesift_ocr/formula/
decoder.rs
