diff --git a/src/lib.rs b/src/lib.rs index 6b83a7d..6962c90 100644 --- a/src/lib.rs +++ b/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, 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, 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) -> 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, + ) -> 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,12 +247,20 @@ 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| { + 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). @@ -239,13 +268,17 @@ impl 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| 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) { + 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. diff --git a/src/post_processors.rs b/src/post_processors.rs index 070e4bd..2d58cc3 100644 --- a/src/post_processors.rs +++ b/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) -> Vec { + pub(crate) fn apply_single(&self, encoded: Vec) -> Vec { 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) { + 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, add_special_tokens: bool) -> Vec { + pub(crate) fn post_process_single( + &self, + encoded: Vec, + add_special_tokens: bool, + ) -> Vec { 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, + ) { + 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)] diff --git a/src/pre_tokenized.rs b/src/pre_tokenized.rs index 97ee5ac..705de0b 100644 --- a/src/pre_tokenized.rs +++ b/src/pre_tokenized.rs @@ -111,11 +111,22 @@ impl PreTokenizedString { /// directly. When there are enough splits, chunks are processed in /// parallel. pub fn tokenize(&self, tokenize_fn: F) -> Result, String> + where + F: Fn(&str, &mut Vec) -> 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(&self, tokenize_fn: F, out: &mut Vec) -> Result<(), String> where F: Fn(&str, &mut Vec) -> 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,11 +152,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(()) }) } @@ -153,13 +164,25 @@ impl PreTokenizedString { /// of splits, allowing it to amortize per-call overhead (e.g. thread-local /// cache access) across the entire chunk. pub fn tokenize_batched(&self, tokenize_fn: F) -> Result, String> + where + F: Fn(&str, &[Split], &mut Vec) -> 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(&self, tokenize_fn: F, out: &mut Vec) -> Result<(), String> where F: Fn(&str, &[Split], &mut Vec) -> 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) -> 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(&self, tokenize_fn: &F) -> Result, String> + /// Sequential tokenization, appending to `out`. + fn tokenize_sequential_into(&self, tokenize_fn: &F, out: &mut Vec) -> Result<(), String> where F: Fn(&str, &mut Vec) -> 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(()) } }