Add encode_into() API to avoid allocations on the encoding hot path

src/lib.rs

@@ -200,7 +200,9 @@ impl Tokenizer {
     /// Run the full encoding pipeline: split added tokens, normalize,
     /// pre-tokenize, tokenize and post-process the input string.
     pub fn encode(&self, input: &str) -> Result<Vec<u32>, Error> {
-        self.encode_with_special_tokens(input, false)
+        let mut ids = Vec::new();
+        self.encode_into(input, &mut ids)?;
+        Ok(ids)
     }
 
     /// Run the full encoding pipeline with control over special token insertion.
@@ -212,12 +214,31 @@ impl Tokenizer {
         input: &str,
         add_special_tokens: bool,
     ) -> Result<Vec<u32>, Error> {
+        let mut ids = Vec::new();
+        self.encode_with_special_tokens_into(input, add_special_tokens, &mut ids)?;
+        Ok(ids)
+    }
+
+    /// Like [`encode`](Self::encode), but appends token IDs to the
+    /// caller-supplied `out` buffer instead of allocating a new `Vec`.
+    pub fn encode_into(&self, input: &str, out: &mut Vec<u32>) -> Result<(), Error> {
+        self.encode_with_special_tokens_into(input, false, out)
+    }
+
+    /// Like [`encode_with_special_tokens`](Self::encode_with_special_tokens),
+    /// but appends token IDs to the caller-supplied `out` buffer instead of
+    /// allocating a new `Vec`.
+    pub fn encode_with_special_tokens_into(
+        &self,
+        input: &str,
+        add_special_tokens: bool,
+        out: &mut Vec<u32>,
+    ) -> Result<(), Error> {
         if input.is_empty() {
-            return if add_special_tokens {
-                Ok(self.post_process(Vec::new(), true))
-            } else {
-                Ok(Vec::new())
-            };
+            if add_special_tokens {
+                self.post_process_into(&[], add_special_tokens, out);
+            }
+            return Ok(());
         }
 
         // 1. Split on added tokens + normalize into a single buffer.
@@ -226,26 +247,38 @@ impl Tokenizer {
         // Fused path: run only Split, then batch-tokenize with inline ByteLevel.
         if let Some(ref split) = self.split_only {
             split.pre_tokenize(&mut pts)?;
-            let ids = pts
-                .tokenize_batched(|buf, splits, out| {
-                    self.model.tokenize_batch_fused(buf, splits, out)
-                })
-                .map_err(Error::Model)?;
-            return Ok(self.post_process(ids, add_special_tokens));
+
+            let tok_fn = |buf: &str, splits: &[crate::pre_tokenized::Split], o: &mut Vec<u32>| {
+                self.model.tokenize_batch_fused(buf, splits, o)
+            };
+            if self.needs_post_process(add_special_tokens) {
+                let mut ids = Vec::new();
+                pts.tokenize_batched_into(tok_fn, &mut ids)
+                    .map_err(Error::Model)?;
+                self.post_process_into(&ids, add_special_tokens, out);
+            } else {
+                pts.tokenize_batched_into(tok_fn, out)
+                    .map_err(Error::Model)?;
+            }
+            return Ok(());
         }
 
         // 2. Pre-tokenize (refine splits in place).
         if let Some(ref pt) = self.pre_tokenizer {
             pt.pre_tokenize(&mut pts)?;
         }
 
-        // 3. Tokenize each text split with the model.
-        let ids = pts
-            .tokenize(|text, out| self.model.tokenize_into(text, out))
-            .map_err(Error::Model)?;
+        // 3. Tokenize each text split with the model + 4. Post-process.
+        let tok_fn = |text: &str, o: &mut Vec<u32>| self.model.tokenize_into(text, o);
+        if self.needs_post_process(add_special_tokens) {
+            let mut ids = Vec::new();
+            pts.tokenize_into(tok_fn, &mut ids).map_err(Error::Model)?;
+            self.post_process_into(&ids, add_special_tokens, out);
+        } else {
+            pts.tokenize_into(tok_fn, out).map_err(Error::Model)?;
+        }
 
-        // 4. Post-process.
-        Ok(self.post_process(ids, add_special_tokens))
+        Ok(())
     }
 
     /// Encode a batch of inputs.
@@ -273,6 +306,25 @@ impl Tokenizer {
         }
     }
 
+    /// Append post-processed token IDs to `out`.
+    fn post_process_into(&self, ids: &[u32], add_special_tokens: bool, out: &mut Vec<u32>) {
+        match &self.post_processor {
+            Some(pp) => pp.post_process_single_into(ids, add_special_tokens, out),
+            None => out.extend_from_slice(ids),
+        }
+    }
+
+    /// Returns `true` when post-processing would actually modify the token
+    /// stream, meaning we need an intermediate buffer before writing to `out`.
+    fn needs_post_process(&self, add_special_tokens: bool) -> bool {
+        add_special_tokens
+            && self.post_processor.is_some()
+            && !matches!(
+                &self.post_processor,
+                Some(post_processors::PostProcessor::ByteLevel)
+            )
+    }
+
     // ── Decoding ─────────────────────────────────────────────────────
 
     /// Decode token IDs back into text.

src/post_processors.rs

@@ -108,25 +108,32 @@ impl TemplateProcessing {
 
     /// Apply the single-sequence template, inserting special token IDs
     /// around the encoded sequence.
-    pub fn apply_single(&self, encoded: Vec<u32>) -> Vec<u32> {
+    pub(crate) fn apply_single(&self, encoded: Vec<u32>) -> Vec<u32> {
         let mut result = Vec::with_capacity(encoded.len() + 4);
+        self.apply_single_into(&encoded, &mut result);
+        result
+    }
+
+    /// Like [`apply_single`](Self::apply_single), but appends the result to
+    /// the caller-supplied `out` buffer.
+    pub(crate) fn apply_single_into(&self, encoded: &[u32], out: &mut Vec<u32>) {
+        out.reserve(encoded.len() + 4);
         for piece in &self.single {
             match piece {
                 TemplatePiece::Sequence {
                     id: SequenceId::A, ..
                 } => {
-                    result.extend_from_slice(&encoded);
+                    out.extend_from_slice(encoded);
                 }
                 TemplatePiece::SpecialToken { id, .. } => {
                     if let Some(ids) = self.special_tokens.get(id) {
-                        result.extend_from_slice(ids);
+                        out.extend_from_slice(ids);
                     }
                 }
                 // Sequence B in a single template is ignored.
                 _ => {}
             }
         }
-        result
     }
 }
 
@@ -182,7 +189,11 @@ impl PostProcessor {
     ///
     /// Only has an effect when `add_special_tokens` is true and the processor
     /// adds special tokens (e.g. `TemplateProcessing`).
-    pub fn post_process_single(&self, encoded: Vec<u32>, add_special_tokens: bool) -> Vec<u32> {
+    pub(crate) fn post_process_single(
+        &self,
+        encoded: Vec<u32>,
+        add_special_tokens: bool,
+    ) -> Vec<u32> {
         if !add_special_tokens {
             return encoded;
         }
@@ -194,6 +205,36 @@ impl PostProcessor {
             }),
         }
     }
+
+    /// Like [`post_process_single`](Self::post_process_single), but appends
+    /// the result to the caller-supplied `out` buffer.
+    ///
+    /// `encoded` is the slice of token IDs produced by tokenization (not yet
+    /// in `out`). The post-processor interleaves special tokens and content,
+    /// appending everything to `out`.
+    pub(crate) fn post_process_single_into(
+        &self,
+        encoded: &[u32],
+        add_special_tokens: bool,
+        out: &mut Vec<u32>,
+    ) {
+        if !add_special_tokens {
+            out.extend_from_slice(encoded);
+            return;
+        }
+        match self {
+            Self::ByteLevel => out.extend_from_slice(encoded),
+            Self::TemplateProcessing(tp) => tp.apply_single_into(encoded, out),
+            Self::Sequence(steps) => {
+                // Sequence is rare; fall back to allocating for intermediate
+                // steps and extend `out` with the final result.
+                let ids = steps.iter().fold(encoded.to_vec(), |acc, step| {
+                    step.post_process_single(acc, add_special_tokens)
+                });
+                out.extend_from_slice(&ids);
+            }
+        }
+    }
 }
 
 #[cfg(test)]

src/pre_tokenized.rs

@@ -111,11 +111,22 @@ impl PreTokenizedString {
     /// directly. When there are enough splits, chunks are processed in
     /// parallel.
     pub fn tokenize<F>(&self, tokenize_fn: F) -> Result<Vec<u32>, String>
+    where
+        F: Fn(&str, &mut Vec<u32>) -> Result<(), String> + Sync,
+    {
+        let mut ids = Vec::new();
+        self.tokenize_into(tokenize_fn, &mut ids)?;
+        Ok(ids)
+    }
+
+    /// Like [`tokenize`](Self::tokenize), but appends token IDs to the
+    /// caller-supplied `out` buffer instead of allocating a new `Vec`.
+    pub fn tokenize_into<F>(&self, tokenize_fn: F, out: &mut Vec<u32>) -> Result<(), String>
     where
         F: Fn(&str, &mut Vec<u32>) -> Result<(), String> + Sync,
     {
         if self.splits.len() < PARALLEL_THRESHOLD {
-            return self.tokenize_sequential(&tokenize_fn);
+            return self.tokenize_sequential_into(&tokenize_fn, out);
         }
 
         let pool = bpe_pool();
@@ -141,25 +152,37 @@ impl PreTokenizedString {
 
             let chunks = chunk_results?;
             let total: usize = chunks.iter().map(Vec::len).sum();
-            let mut ids = Vec::with_capacity(total);
+            out.reserve(total);
             for chunk_ids in chunks {
-                ids.extend(chunk_ids);
+                out.extend(chunk_ids);
             }
-            Ok(ids)
+            Ok(())
         })
     }
 
     /// Batched tokenization: the callback receives the full buffer and a chunk
     /// of splits, allowing it to amortize per-call overhead (e.g. thread-local
     /// cache access) across the entire chunk.
     pub fn tokenize_batched<F>(&self, tokenize_fn: F) -> Result<Vec<u32>, String>
+    where
+        F: Fn(&str, &[Split], &mut Vec<u32>) -> Result<(), String> + Sync,
+    {
+        let mut ids = Vec::new();
+        self.tokenize_batched_into(tokenize_fn, &mut ids)?;
+        Ok(ids)
+    }
+
+    /// Like [`tokenize_batched`](Self::tokenize_batched), but appends token
+    /// IDs to the caller-supplied `out` buffer instead of allocating a new
+    /// `Vec`.
+    pub fn tokenize_batched_into<F>(&self, tokenize_fn: F, out: &mut Vec<u32>) -> Result<(), String>
     where
         F: Fn(&str, &[Split], &mut Vec<u32>) -> Result<(), String> + Sync,
     {
         if self.splits.len() < PARALLEL_THRESHOLD {
-            let mut ids = Vec::with_capacity(self.splits.len() * 2);
-            tokenize_fn(&self.buffer, &self.splits, &mut ids)?;
-            return Ok(ids);
+            out.reserve(self.splits.len() * 2);
+            tokenize_fn(&self.buffer, &self.splits, out)?;
+            return Ok(());
         }
 
         let pool = bpe_pool();
@@ -178,11 +201,11 @@ impl PreTokenizedString {
 
             let chunks = chunk_results?;
             let total: usize = chunks.iter().map(Vec::len).sum();
-            let mut ids = Vec::with_capacity(total);
+            out.reserve(total);
             for chunk_ids in chunks {
-                ids.extend(chunk_ids);
+                out.extend(chunk_ids);
             }
-            Ok(ids)
+            Ok(())
         })
     }
 
@@ -191,26 +214,28 @@ impl PreTokenizedString {
     where
         F: Fn(&str, &mut Vec<u32>) -> Result<(), String>,
     {
-        self.tokenize_sequential(&tokenize_fn)
+        let mut ids = Vec::new();
+        self.tokenize_sequential_into(&tokenize_fn, &mut ids)?;
+        Ok(ids)
     }
 
-    /// Sequential tokenization (used for small inputs).
-    fn tokenize_sequential<F>(&self, tokenize_fn: &F) -> Result<Vec<u32>, String>
+    /// Sequential tokenization, appending to `out`.
+    fn tokenize_sequential_into<F>(&self, tokenize_fn: &F, out: &mut Vec<u32>) -> Result<(), String>
     where
         F: Fn(&str, &mut Vec<u32>) -> Result<(), String>,
     {
-        let mut ids = Vec::with_capacity(self.splits.len() * 2);
+        out.reserve(self.splits.len() * 2);
         for split in &self.splits {
             if let Some(id) = split.token_id {
-                ids.push(id);
+                out.push(id);
             } else {
                 let text = self.split_text(split);
                 if !text.is_empty() {
-                    tokenize_fn(text, &mut ids)?;
+                    tokenize_fn(text, out)?;
                 }
             }
         }
-        Ok(ids)
+        Ok(())
     }
 }