metrale_server/tokenizer/chat_impl/
vision_pads.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-26: Vision placeholder token ids, and the expansion of one placeholder per image
4//! or video into one per vision-encoder output row.
5//!
6//! Owner: server (tokenizer).
7//! Invariants:
8//! - `fan_out_pads` never fails: a placeholder beyond `pad_counts` stays a single token, and
9//!   counts beyond the last placeholder are unused.
10
11use super::super::ChatTokenizer;
12
13impl ChatTokenizer {
14    /// 2026-09-26: The image placeholder token id.
15    ///
16    /// `declared` is the `image_token_id` from the checkpoint's config (top level or under
17    /// `vision_config`) and wins when non-zero. Otherwise the id is that of `<|image_pad|>`
18    /// when the string encodes to exactly one token, else `None`. With `None`,
19    /// `expand_vision_pads` leaves image placeholders unexpanded and reports nothing.
20    pub fn image_pad_token_id(&self, declared: u32) -> Option<u32> {
21        if declared != 0 {
22            return Some(declared);
23        }
24        self.encode("<|image_pad|>")
25            .ok()
26            .and_then(|ids| if ids.len() == 1 { Some(ids[0]) } else { None })
27    }
28
29    /// 2026-09-26: The video placeholder token id, with the same rule as
30    /// `image_pad_token_id`: `declared` when non-zero, else `<|video_pad|>` when it encodes
31    /// to one token, else `None`.
32    pub fn video_pad_token_id(&self, declared: u32) -> Option<u32> {
33        if declared != 0 {
34            return Some(declared);
35        }
36        self.encode("<|video_pad|>")
37            .ok()
38            .and_then(|ids| if ids.len() == 1 { Some(ids[0]) } else { None })
39    }
40
41    /// 2026-09-26: Expand the image and video placeholders of a rendered token sequence, the
42    /// i-th placeholder to `pad_counts[i]` copies (at least one), so the prompt holds one
43    /// placeholder per vision-encoder output row. The Qwen3.6 templates in
44    /// test_data/chat_templates write one placeholder per image or video. Returns `tokens`
45    /// unchanged when every count is at most 1; see `fan_out_pads` for mismatched counts.
46    pub fn expand_vision_pads(
47        &self,
48        tokens: Vec<u32>,
49        pad_counts: &[usize],
50        declared: (u32, u32),
51    ) -> Vec<u32> {
52        if pad_counts.is_empty() || pad_counts.iter().all(|&c| c <= 1) {
53            return tokens;
54        }
55        fan_out_pads(
56            tokens,
57            pad_counts,
58            self.image_pad_token_id(declared.0),
59            self.video_pad_token_id(declared.1),
60        )
61    }
62}
63
64/// 2026-09-26: The expansion over explicit ids, testable without a tokenizer.
65///
66/// `pad_counts[i]` is the number of tokens the i-th media item, image or video, produces.
67/// Placeholders take counts left to right, one each; a placeholder with no count left stays
68/// one token. Returns `tokens` unchanged when both ids are `None`.
69pub(crate) fn fan_out_pads(
70    tokens: Vec<u32>,
71    pad_counts: &[usize],
72    image_pad: Option<u32>,
73    video_pad: Option<u32>,
74) -> Vec<u32> {
75    if image_pad.is_none() && video_pad.is_none() {
76        return tokens;
77    }
78    let extra: usize = pad_counts.iter().map(|c| c.saturating_sub(1)).sum();
79    let mut out = Vec::with_capacity(tokens.len() + extra);
80    let mut img_idx = 0usize;
81    for t in tokens {
82        if Some(t) == image_pad || Some(t) == video_pad {
83            let count = pad_counts.get(img_idx).copied().unwrap_or(1).max(1);
84            for _ in 0..count {
85                out.push(t);
86            }
87            img_idx += 1;
88        } else {
89            out.push(t);
90        }
91    }
92    out
93}
94
95#[cfg(test)]
96mod tests {
97    use super::fan_out_pads;
98
99    /// 2026-09-26: 154854 is the `image_token_id` in the GLM-5.3 config fixture
100    /// (crates/model-engine/tests/fixtures/glm53-nvfp4-9e0d74e3-config.json). With that id
101    /// the placeholder expands; with no id nothing changes.
102    #[test]
103    fn a_declared_id_expands_where_a_literal_probe_would_not_have() {
104        let glm_image = 154_854u32;
105        let tokens = vec![1, 154_830, glm_image, 154_831, 2];
106        let out = fan_out_pads(tokens.clone(), &[256], Some(glm_image), None);
107        assert_eq!(out.len(), tokens.len() + 255);
108        assert_eq!(out.iter().filter(|&&t| t == glm_image).count(), 256);
109        assert_eq!(out[1], 154_830);
110        assert_eq!(*out.last().unwrap(), 2);
111        assert_eq!(out[out.len() - 2], 154_831);
112
113        assert_eq!(fan_out_pads(tokens.clone(), &[256], None, None), tokens);
114    }
115
116    /// 2026-09-26: Images and videos take counts from one left-to-right cursor.
117    #[test]
118    fn counts_are_consumed_in_order_across_both_modalities() {
119        let (img, vid) = (100u32, 101u32);
120        let out = fan_out_pads(vec![img, 7, vid, img], &[2, 3, 1], Some(img), Some(vid));
121        assert_eq!(out, vec![img, img, 7, vid, vid, vid, img]);
122    }
123
124    /// 2026-09-26: More placeholders than counts: the surplus stays one token each.
125    #[test]
126    fn a_short_count_list_does_not_reuse_the_last_count() {
127        let img = 100u32;
128        let out = fan_out_pads(vec![img, img], &[3], Some(img), None);
129        assert_eq!(out, vec![img, img, img, img]);
130    }
131}
metrale_server/tokenizer/chat_impl/vision_pads.rs

metrale_server/tokenizer/chat_impl/
vision_pads.rs
