spark_server/tokenizer/chat_impl/vision_pads.rs
1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! Vision placeholder tokens and the one-pad-per-image → one-pad-per-patch
4//! fan-out.
5//!
6//! Split from `chat_impl.rs` for the 500-LoC cap, and it earns its own file:
7//! this is where the prompt's image placeholders and the vision encoder's
8//! output rows are made to be the same count, and a disagreement between them
9//! is silent in both directions.
10
11use super::super::ChatTokenizer;
12
13impl ChatTokenizer {
14 /// The image placeholder token id.
15 ///
16 /// `declared` is the checkpoint's own `vision_config.image_token_id`, and
17 /// it WINS when non-zero. The literal probe below is the fallback for a
18 /// checkpoint that declares nothing.
19 ///
20 /// 🔴 The probe alone is not enough, and the failure is silent. It looks
21 /// for Qwen's `<|image_pad|>`; GLM-5.3 names the same token `<|image|>`
22 /// (154854), so on GLM the probe returns `None`, `expand_vision_pads`
23 /// becomes a no-op, and the prompt keeps ONE image token where the encoder
24 /// produced 256 embedding rows. Nothing errors: the splice overwrites the
25 /// single position, the remaining 255 rows are never read, and the model
26 /// answers about a 1-token image. The id has been in `VisionConfig` the
27 /// whole time — it just was not asked for.
28 pub fn image_pad_token_id(&self, declared: u32) -> Option<u32> {
29 if declared != 0 {
30 return Some(declared);
31 }
32 self.encode("<|image_pad|>")
33 .ok()
34 .and_then(|ids| if ids.len() == 1 { Some(ids[0]) } else { None })
35 }
36
37 /// `<|video_pad|>`, the temporal sibling, with the same declared-wins rule.
38 /// `None` on a tokenizer without it — every text-only model, and any VL
39 /// model that predates video.
40 pub fn video_pad_token_id(&self, declared: u32) -> Option<u32> {
41 if declared != 0 {
42 return Some(declared);
43 }
44 self.encode("<|video_pad|>")
45 .ok()
46 .and_then(|ids| if ids.len() == 1 { Some(ids[0]) } else { None })
47 }
48
49 /// Post-process a rendered token sequence to expand `<|image_pad|>`
50 /// placeholders. The Qwen3-VL / Qwen3.6 chat template emits exactly one
51 /// `<|image_pad|>` per image, but the vision encoder produces
52 /// `grid_h * grid_w` patches per image. At embed-injection time the
53 /// server expects one pad token per patch so each patch's embedding
54 /// lands at the right hidden-state position — this helper does the
55 /// fan-out.
56 ///
57 /// `pad_counts[i]` is the number of patches the i-th image produces.
58 /// Extra or missing `<|image_pad|>` occurrences (vs `pad_counts.len()`)
59 /// pass through unchanged, matching counts are replicated in place.
60 pub fn expand_vision_pads(
61 &self,
62 tokens: Vec<u32>,
63 pad_counts: &[usize],
64 declared: (u32, u32),
65 ) -> Vec<u32> {
66 if pad_counts.is_empty() || pad_counts.iter().all(|&c| c <= 1) {
67 return tokens;
68 }
69 fan_out_pads(
70 tokens,
71 pad_counts,
72 self.image_pad_token_id(declared.0),
73 self.video_pad_token_id(declared.1),
74 )
75 }
76}
77
78/// The fan-out itself, as a free function over EXPLICIT ids so it can be
79/// tested without a tokenizer.
80///
81/// `pad_counts[i]` is the number of merged tokens the i-th media item
82/// produces. Occurrences are replaced left to right, one count each; a
83/// mismatch in either direction passes the extra tokens through unchanged
84/// rather than guessing.
85pub(crate) fn fan_out_pads(
86 tokens: Vec<u32>,
87 pad_counts: &[usize],
88 image_pad: Option<u32>,
89 video_pad: Option<u32>,
90) -> Vec<u32> {
91 if image_pad.is_none() && video_pad.is_none() {
92 return tokens;
93 }
94 let extra: usize = pad_counts.iter().map(|c| c.saturating_sub(1)).sum();
95 let mut out = Vec::with_capacity(tokens.len() + extra);
96 let mut img_idx = 0usize;
97 for t in tokens {
98 if Some(t) == image_pad || Some(t) == video_pad {
99 let count = pad_counts.get(img_idx).copied().unwrap_or(1).max(1);
100 for _ in 0..count {
101 out.push(t);
102 }
103 img_idx += 1;
104 } else {
105 out.push(t);
106 }
107 }
108 out
109}
110
111#[cfg(test)]
112mod tests {
113 use super::fan_out_pads;
114
115 /// GLM-5.3's id (154854) against a tokenizer that has no `<|image_pad|>`.
116 /// This is the case the literal-probe could not serve: without the
117 /// declared id nothing expands, the prompt keeps ONE image position for
118 /// 256 encoder rows, and the model answers about a 1-token image without
119 /// anything erroring.
120 #[test]
121 fn a_declared_id_expands_where_a_literal_probe_would_not_have() {
122 let glm_image = 154_854u32;
123 let tokens = vec![1, 154_830, glm_image, 154_831, 2];
124 let out = fan_out_pads(tokens.clone(), &[256], Some(glm_image), None);
125 assert_eq!(out.len(), tokens.len() + 255);
126 assert_eq!(out.iter().filter(|&&t| t == glm_image).count(), 256);
127 // The surrounding begin/end-of-image markers are untouched and still
128 // bracket the run.
129 assert_eq!(out[1], 154_830);
130 assert_eq!(*out.last().unwrap(), 2);
131 assert_eq!(out[out.len() - 2], 154_831);
132
133 // Without an id — the pre-fix behaviour on GLM — nothing moves.
134 assert_eq!(fan_out_pads(tokens.clone(), &[256], None, None), tokens);
135 }
136
137 /// Multiple items consume their counts in order, and images and videos
138 /// share one left-to-right cursor.
139 #[test]
140 fn counts_are_consumed_in_order_across_both_modalities() {
141 let (img, vid) = (100u32, 101u32);
142 let out = fan_out_pads(vec![img, 7, vid, img], &[2, 3, 1], Some(img), Some(vid));
143 assert_eq!(out, vec![img, img, 7, vid, vid, vid, img]);
144 }
145
146 /// More placeholders than counts: the surplus passes through as one token
147 /// each rather than reusing the last count.
148 #[test]
149 fn a_short_count_list_does_not_reuse_the_last_count() {
150 let img = 100u32;
151 let out = fan_out_pads(vec![img, img], &[3], Some(img), None);
152 assert_eq!(out, vec![img, img, img, img]);
153 }
154}