tpt_tokenizer_core/
encoding.rs

1//! Batch encoding, special-token insertion, padding, and truncation.
2//!
3//! These conveniences layer on top of the raw [`Tokenizer::encode`] output.
4//! They are scheme-agnostic (they work for both BPE and WordPiece) and live in
5//! the shared [`Tokenizer`] trait as provided methods, so every tokenizer gets
6//! them for free.
7
8use alloc::vec;
9use alloc::vec::Vec;
10
11use crate::error::TokenizerError;
12use crate::tokenizer::{TokenId, Tokenizer};
13
14/// How a sequence (or batch of sequences) should be padded to a common length.
15#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
16pub enum Padding {
17    /// Do not pad. Sequences keep their natural length.
18    #[default]
19    None,
20    /// Pad every sequence to a fixed length. Sequences already at least this
21    /// long are left unchanged (truncation is controlled separately).
22    Fixed(usize),
23    /// Pad every sequence in a batch to the length of the longest sequence.
24    /// Behaves like [`Padding::None`] for a single
25    /// [`encode_with`](TokenizerExt::encode_with) call.
26    Longest,
27}
28
29/// How an over-long sequence should be shortened.
30#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
31pub enum Truncation {
32    /// Never truncate.
33    #[default]
34    None,
35    /// Truncate to at most `n` ids, *including* any special tokens added by
36    /// [`EncodeConfig::add_special_tokens`].
37    Fixed(usize),
38}
39
40/// Configuration for [`encode_with`](TokenizerExt::encode_with) /
41/// [`encode_batch`](TokenizerExt::encode_batch).
42///
43/// Build it fluently:
44///
45/// ```
46/// use tpt_tokenizer_core::encoding::{EncodeConfig, Padding, Truncation};
47///
48/// let cfg = EncodeConfig::new()
49///     .with_bos(1)
50///     .with_eos(2)
51///     .with_pad(0)
52///     .padding(Padding::Longest)
53///     .truncation(Truncation::Fixed(128));
54/// ```
55#[derive(Debug, Clone, Default)]
56pub struct EncodeConfig {
57    bos: Option<TokenId>,
58    eos: Option<TokenId>,
59    pad: Option<TokenId>,
60    padding: Padding,
61    truncation: Truncation,
62    add_special_tokens: bool,
63}
64
65impl EncodeConfig {
66    /// A config that adds no special tokens and applies no padding or
67    /// truncation — equivalent to a plain [`Tokenizer::encode`], but wrapped in
68    /// an [`Encoding`] with an attention mask.
69    #[must_use]
70    pub fn new() -> Self {
71        Self::default()
72    }
73
74    /// Set the beginning-of-sequence token id to prepend (enables
75    /// special-token insertion).
76    #[must_use]
77    pub fn with_bos(mut self, id: TokenId) -> Self {
78        self.bos = Some(id);
79        self.add_special_tokens = true;
80        self
81    }
82
83    /// Set the end-of-sequence token id to append (enables special-token
84    /// insertion).
85    #[must_use]
86    pub fn with_eos(mut self, id: TokenId) -> Self {
87        self.eos = Some(id);
88        self.add_special_tokens = true;
89        self
90    }
91
92    /// Set the padding token id. Required for any [`Padding`] other than
93    /// [`Padding::None`].
94    #[must_use]
95    pub fn with_pad(mut self, id: TokenId) -> Self {
96        self.pad = Some(id);
97        self
98    }
99
100    /// Enable or disable BOS/EOS insertion without changing the ids.
101    #[must_use]
102    pub fn add_special_tokens(mut self, yes: bool) -> Self {
103        self.add_special_tokens = yes;
104        self
105    }
106
107    /// Set the padding strategy.
108    #[must_use]
109    pub fn padding(mut self, padding: Padding) -> Self {
110        self.padding = padding;
111        self
112    }
113
114    /// Set the truncation strategy.
115    #[must_use]
116    pub fn truncation(mut self, truncation: Truncation) -> Self {
117        self.truncation = truncation;
118        self
119    }
120
121    /// Number of special tokens this config would add to a sequence.
122    fn num_special(&self) -> usize {
123        if !self.add_special_tokens {
124            return 0;
125        }
126        usize::from(self.bos.is_some()) + usize::from(self.eos.is_some())
127    }
128}
129
130/// The result of encoding a single sequence with an [`EncodeConfig`].
131#[derive(Debug, Clone, PartialEq, Eq)]
132pub struct Encoding {
133    /// The token ids, including any special/padding tokens.
134    pub ids: Vec<TokenId>,
135    /// A mask that is `1` for real (or special) tokens and `0` for padding.
136    pub attention_mask: Vec<u32>,
137}
138
139impl Encoding {
140    /// Length of the (possibly padded) sequence.
141    #[must_use]
142    pub fn len(&self) -> usize {
143        self.ids.len()
144    }
145
146    /// Whether the encoding is empty.
147    #[must_use]
148    pub fn is_empty(&self) -> bool {
149        self.ids.is_empty()
150    }
151}
152
153/// Assembles a single [`Encoding`] from raw ids: truncate content to leave room
154/// for specials, insert BOS/EOS, then build the attention mask. Padding is *not*
155/// applied here (it depends on the batch), so all mask entries are `1`.
156pub(crate) fn assemble(mut ids: Vec<TokenId>, cfg: &EncodeConfig) -> Encoding {
157    if let Truncation::Fixed(max) = cfg.truncation {
158        let room = max.saturating_sub(cfg.num_special());
159        if ids.len() > room {
160            ids.truncate(room);
161        }
162    }
163
164    let mut out = Vec::with_capacity(ids.len() + cfg.num_special());
165    if cfg.add_special_tokens {
166        if let Some(bos) = cfg.bos {
167            out.push(bos);
168        }
169    }
170    out.extend(ids);
171    if cfg.add_special_tokens {
172        if let Some(eos) = cfg.eos {
173            out.push(eos);
174        }
175    }
176
177    let attention_mask = vec![1u32; out.len()];
178    Encoding {
179        ids: out,
180        attention_mask,
181    }
182}
183
184/// Pads `enc` up to `target` ids using `pad_id`, extending the attention mask
185/// with zeros. A no-op if already at least `target` long.
186pub(crate) fn pad_to(enc: &mut Encoding, target: usize, pad_id: TokenId) {
187    if enc.ids.len() >= target {
188        return;
189    }
190    let deficit = target - enc.ids.len();
191    enc.ids.extend(core::iter::repeat(pad_id).take(deficit));
192    enc.attention_mask
193        .extend(core::iter::repeat(0u32).take(deficit));
194}
195
196/// Applies the batch-level padding strategy across `encodings` in place.
197///
198/// # Errors
199/// Returns [`TokenizerError::MalformedFile`] if padding is requested but no pad
200/// token id was configured.
201pub(crate) fn apply_padding(
202    encodings: &mut [Encoding],
203    cfg: &EncodeConfig,
204) -> Result<(), TokenizerError> {
205    let target = match cfg.padding {
206        Padding::None => return Ok(()),
207        Padding::Fixed(n) => n,
208        Padding::Longest => encodings.iter().map(Encoding::len).max().unwrap_or(0),
209    };
210    let Some(pad_id) = cfg.pad else {
211        return Err(TokenizerError::MalformedFile(alloc::string::String::from(
212            "padding requested without a pad token id (call EncodeConfig::with_pad)",
213        )));
214    };
215    for enc in encodings.iter_mut() {
216        pad_to(enc, target, pad_id);
217    }
218    Ok(())
219}
220
221/// Provided-method extensions shared by every [`Tokenizer`].
222impl<T: Tokenizer + ?Sized> TokenizerExt for T {}
223
224/// Higher-level encoding conveniences (special tokens, padding, truncation,
225/// batching) available on every [`Tokenizer`].
226pub trait TokenizerExt: Tokenizer {
227    /// Encode a single sequence, applying special-token insertion, truncation
228    /// and (fixed) padding from `cfg`.
229    ///
230    /// [`Padding::Longest`] has no effect on a single sequence — use
231    /// [`encode_batch`](TokenizerExt::encode_batch) for that.
232    ///
233    /// # Errors
234    /// Propagates [`Tokenizer::encode`] errors, or returns
235    /// [`TokenizerError::MalformedFile`] if fixed padding is requested without a
236    /// pad token id.
237    fn encode_with(&self, text: &str, cfg: &EncodeConfig) -> Result<Encoding, TokenizerError> {
238        let ids = self.encode(text)?;
239        let mut enc = assemble(ids, cfg);
240        if let Padding::Fixed(n) = cfg.padding {
241            let Some(pad_id) = cfg.pad else {
242                return Err(TokenizerError::MalformedFile(alloc::string::String::from(
243                    "padding requested without a pad token id (call EncodeConfig::with_pad)",
244                )));
245            };
246            pad_to(&mut enc, n, pad_id);
247        }
248        Ok(enc)
249    }
250
251    /// Encode a batch of sequences, applying special tokens and truncation to
252    /// each and then the batch-level [`Padding`] strategy across all of them.
253    ///
254    /// # Errors
255    /// Propagates [`Tokenizer::encode`] errors, or returns
256    /// [`TokenizerError::MalformedFile`] if padding is requested without a pad
257    /// token id.
258    fn encode_batch(
259        &self,
260        texts: &[&str],
261        cfg: &EncodeConfig,
262    ) -> Result<Vec<Encoding>, TokenizerError> {
263        let mut encodings = Vec::with_capacity(texts.len());
264        for text in texts {
265            let ids = self.encode(text)?;
266            encodings.push(assemble(ids, cfg));
267        }
268        apply_padding(&mut encodings, cfg)?;
269        Ok(encodings)
270    }
271}
tpt_tokenizer_core/encoding.rs

tpt_tokenizer_core/
encoding.rs
