Skip to main content

dynamo_tokenizers/
lib.rs

1// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4pub mod basetenkenizer;
5pub mod cache;
6pub mod fastokens;
7pub mod hf;
8pub mod tiktoken;
9
10// TODO: Add tokenizer benchmarks
11// TODO: Enable README.md as a module doc
12// #[doc = include_str!("../README.md")]
13
14use std::hash::{DefaultHasher, Hash, Hasher};
15use std::sync::Arc;
16use std::{fs::File, io::BufReader, ops::Deref, path::Path};
17
18use anyhow::Context as _;
19pub use anyhow::{Error, Result};
20
21pub use basetenkenizer::BasetenTokenizer;
22pub use cache::{CacheTokenUsage, CacheTokenUsageFn, CachedTokenizer, L1CacheStats};
23pub use fastokens::FastTokenizer;
24pub use hf::HuggingFaceTokenizer;
25pub use tiktoken::TikTokenTokenizer;
26pub use traits::DecodeResult;
27
28pub type TokenIdType = u32;
29
30/// A rendered prompt segment with an explicit trust boundary for special tokens.
31#[derive(Debug, Clone, Copy, PartialEq, Eq)]
32pub struct EncodeSegment<'a> {
33    pub text: &'a str,
34    /// Recognize added/control tokens in this trusted renderer output.
35    ///
36    /// Set this to `false` for user, tool, and attribute content so text that
37    /// resembles a control token is encoded as ordinary text.
38    pub allow_special: bool,
39}
40
41impl<'a> EncodeSegment<'a> {
42    pub const fn new(text: &'a str, allow_special: bool) -> Self {
43        Self {
44            text,
45            allow_special,
46        }
47    }
48
49    pub const fn ordinary(text: &'a str) -> Self {
50        Self::new(text, false)
51    }
52
53    pub const fn control(text: &'a str) -> Self {
54        Self::new(text, true)
55    }
56}
57
58/// Represents the type of tokenizer being used
59#[derive(Debug)]
60pub enum TokenizerType {
61    HuggingFace(String),
62    TikToken(String),
63}
64
65/// character offsets in the original text
66pub type Offsets = (usize, usize);
67
68/// Contains the results of tokenizing text: token IDs, string tokens, and their spans
69#[derive(Debug, Clone)]
70pub enum Encoding {
71    /// Hugging Face
72    Hf(Box<tokenizers::tokenizer::Encoding>),
73    /// Sentence Piece
74    Sp(Vec<TokenIdType>),
75}
76
77impl Encoding {
78    pub fn token_ids(&self) -> &[u32] {
79        match self {
80            Encoding::Hf(inner) => inner.get_ids(),
81            Encoding::Sp(inner) => inner,
82        }
83    }
84}
85
86impl Hash for Encoding {
87    fn hash<H: Hasher>(&self, state: &mut H) {
88        self.token_ids().hash(state);
89    }
90}
91
92pub mod traits {
93    use super::*;
94
95    pub trait Encoder: Send + Sync {
96        fn encode(&self, input: &str) -> Result<Encoding>;
97        fn encode_batch(&self, inputs: &[&str]) -> Result<Vec<Encoding>>;
98
99        /// Encode Kimi K3-style renderer segments while preserving trusted
100        /// control-token and untrusted content boundaries.
101        ///
102        /// Each segment controls whether added/control tokens are recognized
103        /// through [`EncodeSegment::allow_special`]. This prevents text in user,
104        /// tool, or attribute content from becoming structural when it happens
105        /// to resemble a control token.
106        ///
107        /// The Baseten backend preserves legacy tiktoken behavior by splitting
108        /// each segment into chunks of at most 400,000 characters and splitting
109        /// whitespace/non-whitespace runs at 25,000 characters. Independent
110        /// chunks are encoded through the Rayon thread pool, then concatenated
111        /// in input order. Tokenizer post-processing is applied once after all
112        /// segment IDs have been joined.
113        ///
114        /// Backends must not implement this by flattening the segments first,
115        /// because that discards the special-token trust boundary.
116        fn encode_segments(&self, _segments: &[EncodeSegment<'_>]) -> Result<Encoding> {
117            Err(Error::msg(
118                "tokenizer backend does not support segmented encoding",
119            ))
120        }
121    }
122
123    /// Result of decoding token IDs to text.
124    ///
125    /// Distinguishes between fully valid UTF-8 output and output that contains
126    /// trailing incomplete multi-byte sequences (represented as U+FFFD).
127    /// This lets callers like `DecodeStream::step()` decide whether to emit or
128    /// buffer without resorting to hardcoded replacement-character string checks.
129    #[derive(Debug, Clone, PartialEq, Eq, strum::EnumIs)]
130    pub enum DecodeResult {
131        /// No trailing incomplete multi-byte sequences (text does not end with U+FFFD).
132        /// Note: the string may still contain *interior* U+FFFD characters from
133        /// mid-stream invalid byte sequences; only trailing status is tracked here.
134        Complete(String),
135        /// The decoded string ends with U+FFFD, indicating incomplete trailing
136        /// multi-byte bytes that may be completed by subsequent tokens.
137        Partial(String),
138    }
139
140    impl DecodeResult {
141        /// Returns a reference to the inner string.
142        pub fn as_str(&self) -> &str {
143            match self {
144                DecodeResult::Complete(s) | DecodeResult::Partial(s) => s,
145            }
146        }
147
148        /// Construct from a decoded string: `Partial` if it ends with U+FFFD, else `Complete`.
149        pub fn from_decoded(text: String) -> Self {
150            if text.ends_with('\u{FFFD}') {
151                DecodeResult::Partial(text)
152            } else {
153                DecodeResult::Complete(text)
154            }
155        }
156    }
157
158    impl From<String> for DecodeResult {
159        fn from(text: String) -> Self {
160            DecodeResult::from_decoded(text)
161        }
162    }
163
164    impl From<DecodeResult> for String {
165        fn from(result: DecodeResult) -> Self {
166            match result {
167                DecodeResult::Complete(s) | DecodeResult::Partial(s) => s,
168            }
169        }
170    }
171
172    /// Implementations must ensure that partial multi-byte sequences produce U+FFFD
173    /// (`\u{FFFD}`) in the output rather than returning `Err`. This is commonly achieved
174    /// via `String::from_utf8_lossy` (tiktoken) or library-internal byte-fallback handling
175    /// (HuggingFace). `DecodeStream::step()` relies on `DecodeResult::Partial` to detect
176    /// incomplete sequences and buffer tokens until the full character arrives.
177    pub trait Decoder: Send + Sync {
178        fn decode(
179            &self,
180            token_ids: &[TokenIdType],
181            skip_special_tokens: bool,
182        ) -> Result<DecodeResult>;
183    }
184
185    pub trait Tokenizer: Encoder + Decoder {
186        /// Validate that this tokenizer can be safely wrapped in the prefix cache.
187        ///
188        /// Implementations must explicitly opt in by returning `Ok(())`, or
189        /// return an error explaining why the tokenizer is incompatible.
190        fn validate_prefix_cache(&self) -> Result<()> {
191            Err(Error::msg("tokenizer does not support prefix caching"))
192        }
193
194        /// Apply construction-time [`TokenizerOptions`].
195        ///
196        /// The default implementation ignores the options — correct for
197        /// tokenizers with no applicable option.
198        fn with_options(self, options: TokenizerOptions) -> Self
199        where
200            Self: Sized,
201        {
202            let _ = options;
203            self
204        }
205        // fn get_vocab_size(&self) -> usize;
206        // fn make_unique_clone(&self) -> Box<dyn Tokenizer>;
207    }
208}
209
210pub fn file_json_field<T: serde::de::DeserializeOwned>(
211    json_file_path: &Path,
212    field_name: &str,
213) -> anyhow::Result<T> {
214    let file = File::open(json_file_path)
215        .with_context(|| format!("Failed to open file: {:?}", json_file_path))?;
216    let reader = BufReader::new(file);
217
218    let json_data: serde_json::Value = serde_json::from_reader(reader)
219        .with_context(|| format!("Failed to parse JSON from file: {:?}", json_file_path))?;
220
221    let map = json_data.as_object().ok_or_else(|| {
222        anyhow::anyhow!("JSON root is not an object in file: {:?}", json_file_path)
223    })?;
224
225    let field_value = map.get(field_name).ok_or_else(|| {
226        anyhow::anyhow!(
227            "Field '{}' not found in JSON file: {:?}",
228            field_name,
229            json_file_path
230        )
231    })?;
232
233    serde_json::from_value(field_value.clone()).with_context(|| {
234        format!(
235            "Failed to deserialize field '{}' (value: {:?}) to the expected type from file: {:?}",
236            field_name, field_value, json_file_path
237        )
238    })
239}
240
241pub fn log_json_err(filename: &str, json: &str, err: &serde_json::Error) {
242    const ERROR_PREFIX: &str = ">>     ";
243
244    if !(err.is_syntax() || err.is_data()) {
245        return;
246    }
247
248    let line = err.line().saturating_sub(1);
249    let column = err.column().saturating_sub(1);
250
251    let json_lines: Vec<&str> = json.lines().collect();
252    if json_lines.is_empty() {
253        tracing::error!("JSON parsing error in {filename}: File is empty.");
254        return;
255    }
256
257    let start_index = line.saturating_sub(2);
258    let end_index = line.saturating_add(3).min(json_lines.len());
259
260    let mut context_lines: Vec<String> = (start_index..end_index)
261        .map(|i| {
262            if i == line {
263                format!("{ERROR_PREFIX}{}", json_lines[i])
264            } else {
265                format!("{:06} {}", i + 1, json_lines[i])
266            }
267        })
268        .collect();
269
270    let col_indicator = "_".to_string().repeat(column + ERROR_PREFIX.len()) + "^";
271    let error_in_context_idx = line - start_index;
272    if error_in_context_idx < context_lines.len() {
273        context_lines.insert(error_in_context_idx + 1, col_indicator);
274    }
275
276    tracing::error!(
277        "JSON parsing error in {filename}: Line {}, column {}:\n{}",
278        err.line(),
279        err.column(),
280        context_lines.join("\n")
281    );
282}
283
284impl Encoding {
285    pub fn get_hash(&self) -> u64 {
286        let mut hasher = DefaultHasher::new();
287        self.hash(&mut hasher);
288        hasher.finish()
289    }
290}
291
292/// Construction options for [`Tokenizer::from_file_with_options`] /
293/// [`create_tokenizer_from_file`], applied to concrete tokenizers via
294/// [`traits::Tokenizer::with_options`].
295#[derive(Debug, Clone, Copy, Default)]
296pub struct TokenizerOptions {
297    /// Ask the tokenizer to add its declared special tokens (e.g. BOS/EOS via
298    /// its post-processor) during `encode`, `encode_batch`, and supported
299    /// `encode_segments` calls.
300    /// Defaults to `false`, the historical behavior.
301    ///
302    /// Applicable to Hugging Face and Baseten tokenizers.
303    pub add_special_tokens: bool,
304}
305
306/// Main tokenizer wrapper that provides a unified interface for different tokenizer implementations
307#[derive(Clone)]
308pub struct Tokenizer(Arc<dyn traits::Tokenizer>);
309
310impl Tokenizer {
311    pub fn from_file(file_path: &str) -> Result<Tokenizer> {
312        Ok(Tokenizer(create_tokenizer_from_file(file_path)?))
313    }
314
315    pub fn from_file_with_options(file_path: &str, options: TokenizerOptions) -> Result<Tokenizer> {
316        Ok(Tokenizer(create_tokenizer_from_file_with_options(
317            file_path, options,
318        )?))
319    }
320
321    /// Create a stateful sequence object for decoding token_ids into text
322    pub fn decode_stream(
323        &self,
324        prompt_token_ids: &[TokenIdType],
325        skip_special_tokens: bool,
326    ) -> DecodeStream {
327        DecodeStream::new(self.0.clone(), prompt_token_ids, skip_special_tokens)
328    }
329}
330
331impl Deref for Tokenizer {
332    type Target = Arc<dyn traits::Tokenizer>;
333
334    fn deref(&self) -> &Self::Target {
335        &self.0
336    }
337}
338
339impl From<Arc<dyn traits::Tokenizer>> for Tokenizer {
340    fn from(tokenizer: Arc<dyn traits::Tokenizer>) -> Self {
341        Tokenizer(tokenizer)
342    }
343}
344
345impl<T> From<Arc<T>> for Tokenizer
346where
347    T: traits::Tokenizer + 'static, // 'static is required to ensure T can be safely put into an Arc
348{
349    fn from(tokenizer: Arc<T>) -> Self {
350        Tokenizer(tokenizer)
351    }
352}
353
354/// Create a tokenizer from a file path to a tokenizer file.
355/// The file extension is used to determine the tokenizer type.
356/// Supported file types are:
357/// - json: HuggingFace tokenizer
358/// - model, tiktoken: tiktoken BPE tokenizer (requires `config.json` with a supported
359///   `model_type` in the same directory; currently: kimi, kimi_k2, kimi_k25, kimi_k3)
360pub fn create_tokenizer_from_file(file_path: &str) -> Result<Arc<dyn traits::Tokenizer>> {
361    create_tokenizer_from_file_with_options(file_path, Default::default())
362}
363
364/// Create a tokenizer from a file path to a tokenizer file with additional tokenizer option.
365/// The file extension is used to determine the tokenizer type.
366/// Supported file types are:
367/// - json: HuggingFace tokenizer
368/// - model, tiktoken: tiktoken BPE tokenizer (requires `config.json` with a supported
369///   `model_type` in the same directory; currently: kimi, kimi_k2, kimi_k25, kimi_k3)
370pub fn create_tokenizer_from_file_with_options(
371    file_path: &str,
372    options: TokenizerOptions,
373) -> Result<Arc<dyn traits::Tokenizer>> {
374    use traits::Tokenizer as _;
375
376    let path = Path::new(file_path);
377    let extension = path
378        .extension()
379        .and_then(std::ffi::OsStr::to_str)
380        .ok_or_else(|| Error::msg("Failed to read file extension".to_string()))?;
381
382    match extension {
383        "json" => {
384            let tokenizer = HuggingFaceTokenizer::from_file(file_path)?.with_options(options);
385            Ok(Arc::new(tokenizer))
386        }
387        "model" | "tiktoken" => {
388            let tokenizer = TikTokenTokenizer::from_file_auto(file_path)?.with_options(options);
389            Ok(Arc::new(tokenizer))
390        }
391        _ => Err(Error::msg(format!(
392            "Unsupported tokenizer file type: .{extension}"
393        ))),
394    }
395}
396
397// With incremental detokenization, we need to consider the final context tokens when handling the initial decode tokens.
398// This is the initial offset from the end of the context that we start decoding from.
399// Both Huggingface TGI and vLLM use this same value.
400// See: https://github.com/huggingface/text-generation-inference/blob/24c2bff65924801ddf90fa24fcc72752d4f45538/server/text_generation_server/models/mamba.py#L169
401// and https://github.com/vllm-project/vllm/blob/da2705198fa19030a25d0bea437f7be6547d47d4/vllm/transformers_utils/detokenizer_utils.py#L51
402const INITIAL_INCREMENTAL_DETOKENIZATION_OFFSET: usize = 5;
403
404/// DecodeStream will keep the state necessary to produce individual chunks of
405/// strings given an input stream of token_ids.
406///
407/// This is necessary because decoding in general cannot achieve that since strings
408/// depend on surrounding ids to provide a valid string. Typically stripping extra spaces.
409pub struct DecodeStream {
410    /// The tokenizer used to decode token_ids
411    tokenizer: Arc<dyn traits::Tokenizer>,
412
413    skip_special_tokens: bool,
414    /// A temporary buffer of the necessary token_ids needed
415    /// to produce valid string chunks.
416    /// This typically contains 3 parts:
417    ///  - read
418    ///  - prefix
419    ///  - rest
420    ///
421    /// Read is the bit necessary to surround the prefix
422    /// so decoding the whole ids produces a valid prefix.
423    /// Prefix is the previously produced string, kept around to trim off of
424    /// the next valid chunk
425    all_token_ids: Vec<u32>,
426
427    prefix_offset: usize,
428
429    read_offset: usize,
430}
431
432impl DecodeStream {
433    pub fn new(
434        tokenizer: Arc<dyn traits::Tokenizer>,
435        prompt_token_ids: &[TokenIdType],
436        skip_special_tokens: bool,
437    ) -> Self {
438        let num_input_tokens = prompt_token_ids.len();
439        let prompt_token_ids = prompt_token_ids.to_vec();
440        Self {
441            tokenizer,
442            skip_special_tokens,
443            all_token_ids: prompt_token_ids,
444            prefix_offset: num_input_tokens
445                .saturating_sub(INITIAL_INCREMENTAL_DETOKENIZATION_OFFSET),
446            read_offset: num_input_tokens,
447        }
448    }
449
450    /// Step appends a token_id to the internal state and tries to produce a text chunk.
451    ///
452    /// Implementation directly copied from Huggingface's TGI:
453    /// https://github.com/huggingface/text-generation-inference/blob/24c2bff65924801ddf90fa24fcc72752d4f45538/server/text_generation_server/models/model.py#L144
454    ///
455    /// Returning `None` means the given id is not enough to produce a chunk.
456    /// This typically happens with `byte_fallback` options where some tokens do not
457    /// represent valid UTF-8, and only follow-up token_ids will help produce
458    /// a valid chunk.
459    pub fn step(&mut self, id: u32) -> Result<Option<String>> {
460        self.all_token_ids.push(id);
461
462        let prefix_text: String = self
463            .tokenizer
464            .decode(
465                &self.all_token_ids[self.prefix_offset..self.read_offset],
466                self.skip_special_tokens,
467            )?
468            .into();
469
470        let new_result = self.tokenizer.decode(
471            &self.all_token_ids[self.prefix_offset..],
472            self.skip_special_tokens,
473        )?;
474
475        let new_text = new_result.as_str();
476        if new_text.len() > prefix_text.len() && !new_result.is_partial() {
477            let emitted = new_text[prefix_text.len()..].to_string();
478
479            self.prefix_offset = self.read_offset;
480            self.read_offset = self.all_token_ids.len();
481
482            Ok(Some(emitted))
483        } else {
484            Ok(None)
485        }
486    }
487}
488
489/// Maintains state for an ongoing sequence of tokens and their decoded text
490pub struct Sequence {
491    /// Encodes text -> token_ids
492    tokenizer: Tokenizer,
493
494    /// The current sequence of token ids
495    token_ids: Vec<TokenIdType>,
496
497    /// The position in the current sequence the last decoded token completed
498    prefix_offset: usize,
499
500    /// Current position in the sequence
501    read_offset: usize,
502}
503
504impl std::fmt::Debug for Sequence {
505    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
506        f.debug_struct("Sequence")
507            .field("tokenizer", &"Arc<dyn Tokenizer>")
508            .field(
509                "token_ids",
510                &format_args!("{}", {
511                    let token_ids = self.token_ids();
512                    if token_ids.len() <= 20 {
513                        format!("{:?}", token_ids)
514                    } else {
515                        let first_ten = &token_ids[..10];
516                        let last_ten = &token_ids[token_ids.len() - 10..];
517                        format!("{:?} ... {:?}", first_ten, last_ten)
518                    }
519                }),
520            )
521            .field("prefix_offset", &self.prefix_offset)
522            .field("read_offset", &self.read_offset)
523            .field("token count", &self.token_ids.len())
524            .finish()
525    }
526}
527
528impl Sequence {
529    pub fn new(tokenizer: Tokenizer) -> Self {
530        Self {
531            tokenizer,
532            token_ids: Vec::new(),
533            prefix_offset: 0,
534            read_offset: 0,
535        }
536    }
537
538    pub fn is_empty(&self) -> bool {
539        self.token_ids.is_empty()
540    }
541
542    pub fn len(&self) -> usize {
543        self.token_ids.len()
544    }
545
546    pub fn clear(&mut self) {
547        self.token_ids.clear();
548        self.prefix_offset = 0;
549        self.read_offset = 0;
550    }
551
552    pub fn append_text(&mut self, input: &str) -> Result<()> {
553        // let tokenizer = self.tokenizer.read().map_err(|err| {
554        //     Error::msg(format!("Failed to acquire read lock on tokenizer: {}", err))
555        // })?;
556
557        let encoding = self.tokenizer.encode(input)?;
558        self.token_ids.extend(encoding.token_ids());
559        Ok(())
560    }
561
562    // Based on
563    // https://github.com/huggingface/text-generation-inference/blob/v0.9.4/server/text_generation_server/models/model.py#L62C9-L62C15
564    // under Apache 2.0 license
565    pub fn append_token_id(&mut self, token_id: TokenIdType) -> Result<String> {
566        self.token_ids.push(token_id);
567        // log::trace!("pushed token_id: {}", token_id);
568
569        let prefix_text: String = self
570            .tokenizer
571            .decode(&self.token_ids[self.prefix_offset..self.read_offset], false)?
572            .into();
573
574        let new_result = self
575            .tokenizer
576            .decode(&self.token_ids[self.prefix_offset..], false)?;
577
578        let new_text = new_result.as_str();
579
580        // if the end character of the previous returned sequence is a multi-byte character
581        // then we can not split the text on that byte offset, so we roll back to the byte offset
582        // of the start of that character
583        let mut prefix_text_len = prefix_text.len();
584        while !new_text.is_char_boundary(prefix_text_len) && prefix_text_len > 0 {
585            prefix_text_len -= 1;
586        }
587        let prefix_text_len = prefix_text_len;
588
589        if new_text.len() > prefix_text.len() {
590            if new_result.is_partial() {
591                return Ok("".to_string());
592            } else {
593                // shift and update the state
594                let new_text = new_text[prefix_text_len..]
595                    .to_string()
596                    .replace('\u{FFFD}', "");
597                self.prefix_offset = self.read_offset;
598                self.read_offset = self.token_ids.len();
599                return Ok(new_text);
600            }
601        }
602
603        Ok("".to_string())
604    }
605
606    pub fn tokenizer(&self) -> Tokenizer {
607        self.tokenizer.clone()
608    }
609
610    pub fn token_ids(&self) -> &[TokenIdType] {
611        &self.token_ids
612    }
613
614    pub fn text(&self) -> Result<String> {
615        // let tokenizer = self.tokenizer.read().map_err(|err| {
616        //     Error::msg(format!("Failed to acquire read lock on tokenizer: {}", err))
617        // })?;
618        Ok(self.tokenizer.decode(&self.token_ids, false)?.into())
619    }
620}
621
622/// The output conditions/values of a SequenceDecoder::add_token_id operation.
623/// Result of decoding a token, indicating whether text was produced or a stop condition was met
624pub enum SequenceDecoderOutput {
625    /// The text for the appended token_id
626    Text(String),
627
628    /// A sequence of token_ids has been partially matched a stop sequence, so the text is held
629    /// until either a match or a divergence
630    Held,
631
632    /// Indicates that a stop sequence has been matched and the decoder is stopped.
633    /// Subsequent calls to append_token_id will return an error
634    Stopped,
635
636    /// Indicates that a stop token_id has been matched and the decoder is stopped.
637    /// Subsequent calls to append_token_id will return an error
638    /// The text for the stop token_id is returned
639    StoppedWithText(String),
640}
641
642/// A Sequence for decoding a stream of token ids into text and detecting stop sequences.
643/// A stop sequence is either a matching token_id or a sequence of texts/strings which match.
644/// Matches happen first at the token-level, then at the sequence-level. Hidden takes precedence
645/// over visible. For example, if you put the same token_id in both `stop_token_ids_visible` and
646/// `stop_token_ids_hidden`, the token_id will be treated as hidden.
647#[derive(Debug)]
648pub struct StopSequenceDecoder {
649    // The current sequence of token ids
650    sequence: Sequence,
651
652    // Stop Tokens - the presence of any one of these should trigger a stop
653    // If found, the text for the matched token will be returned
654    stop_token_ids_visible: Vec<TokenIdType>,
655
656    // Stop Tokens - the presence of any one of these should trigger a stop
657    // If found, the text for the matched token will NOT be returned
658    stop_token_ids_hidden: Vec<TokenIdType>,
659
660    // Stop Words - the presence of any one of these should trigger a stop
661    // If found, the text for the matched token will be returned
662    #[allow(dead_code)]
663    stop_sequences_visible: Vec<String>,
664
665    // Stop Words - the presence of any one of these should trigger a stop
666    // If found, the text for the matched token will NOT be returned
667    stop_sequences_hidden: Vec<String>,
668
669    // If the decoder has observed and returned a stop SequenceDecoderOutput,
670    // futhur calls to append_token_id will return an error
671    stopped: bool,
672
673    // text jail - if a partial stop sequence is being observed, we hold/jail the text
674    // until either the stop sequence is matched or the sequence is reset by a divergence
675    state: String,
676}
677
678impl StopSequenceDecoder {
679    /// Builder object for configurating a StopSequenceDecoder
680    pub fn builder(tokenizer: Tokenizer) -> StopSequenceDecoderBuilder {
681        StopSequenceDecoderBuilder::new(tokenizer)
682    }
683
684    /// Add a token_id to the sequence and return the SequenceDecoderOutput
685    pub fn append_token_id(&mut self, token_id: TokenIdType) -> Result<SequenceDecoderOutput> {
686        if self.stopped {
687            return Err(Error::msg("Decoder is stopped"));
688        }
689
690        // update the sequence
691        let text = self.sequence.append_token_id(token_id)?;
692
693        // append the text to the state
694        self.state.push_str(text.as_str());
695
696        let mut stop: bool = false;
697        let mut visible: bool = false;
698
699        if self.stop_token_ids_visible.contains(&token_id) {
700            stop = true;
701            visible = true;
702        }
703
704        if self.stop_token_ids_hidden.contains(&token_id) {
705            stop = true;
706            visible = false;
707        }
708
709        if stop {
710            self.stopped = true;
711            let state = std::mem::take(&mut self.state);
712            if visible {
713                return Ok(SequenceDecoderOutput::StoppedWithText(state));
714            }
715            return Ok(SequenceDecoderOutput::Stopped);
716        }
717
718        // determine if state matches any of the stop sequences
719        for stop_sequence in self.stop_sequences_hidden.iter() {
720            if stop_sequence.starts_with(&self.state) {
721                if stop_sequence == &self.state {
722                    // on matched stop sequence, we do NOT return the jailed stop sequence
723                    self.stopped = true;
724                    return Ok(SequenceDecoderOutput::Stopped);
725                } else {
726                    return Ok(SequenceDecoderOutput::Held);
727                }
728            }
729        }
730
731        let state = std::mem::take(&mut self.state);
732        Ok(SequenceDecoderOutput::Text(state))
733    }
734
735    pub fn is_empty(&self) -> bool {
736        self.sequence.token_ids.is_empty()
737    }
738
739    pub fn len(&self) -> usize {
740        self.sequence.token_ids.len()
741    }
742
743    pub fn is_complete(&self) -> bool {
744        self.stopped
745    }
746
747    pub fn close(&mut self) {
748        self.stopped = true;
749    }
750}
751
752pub struct StopSequenceDecoderBuilder {
753    tokenizer: Tokenizer,
754    stop_token_ids_visible: Vec<TokenIdType>,
755    stop_token_ids_hidden: Vec<TokenIdType>,
756    stop_sequences_visible: Vec<String>,
757    stop_sequences_hidden: Vec<String>,
758}
759
760impl StopSequenceDecoderBuilder {
761    pub fn new(tokenizer: Tokenizer) -> Self {
762        Self {
763            tokenizer,
764            stop_token_ids_visible: Vec::new(),
765            stop_token_ids_hidden: Vec::new(),
766            stop_sequences_visible: Vec::new(),
767            stop_sequences_hidden: Vec::new(),
768        }
769    }
770
771    /// Adds a visible stop token id to the StopSequenceDecoder
772    pub fn add_stop_token_id_visible(mut self, token_id: TokenIdType) -> Self {
773        self.stop_token_ids_visible.push(token_id);
774        self
775    }
776
777    /// Adds a list of visible stop token ids to the StopSequenceDecoder
778    /// Each token_id is added as for an individual match
779    pub fn add_stop_token_ids_visible(mut self, token_ids: &[TokenIdType]) -> Self {
780        self.stop_token_ids_visible.extend(token_ids);
781        self
782    }
783
784    /// Adds a hidden stop token id to the StopSequenceDecoder
785    pub fn add_stop_token_id_hidden(mut self, token_id: TokenIdType) -> Self {
786        self.stop_token_ids_hidden.push(token_id);
787        self
788    }
789
790    /// Adds a list of hidden stop token ids to the StopSequenceDecoder
791    /// Each token_id is added as for an individual match
792    pub fn add_stop_token_ids_hidden(mut self, token_ids: &[TokenIdType]) -> Self {
793        self.stop_token_ids_hidden.extend(token_ids);
794        self
795    }
796
797    pub fn add_stop_sequence_visible(mut self, text: &str) -> Self {
798        self.stop_sequences_visible.push(text.to_string());
799        self
800    }
801
802    pub fn add_stop_sequences_visible(mut self, strings: &[&str]) -> Self {
803        self.stop_sequences_visible
804            .extend(strings.iter().map(|text| text.to_string()));
805        self
806    }
807
808    pub fn add_stop_sequence_hidden(mut self, text: &str) -> Self {
809        self.stop_sequences_hidden.push(text.to_string());
810        self
811    }
812
813    pub fn add_stop_sequences_hidden(mut self, strings: &[&str]) -> Self {
814        self.stop_sequences_hidden
815            .extend(strings.iter().map(|text| text.to_string()));
816        self
817    }
818
819    pub fn build(self) -> Result<StopSequenceDecoder> {
820        Ok(StopSequenceDecoder {
821            sequence: Sequence::new(self.tokenizer.clone()),
822            stop_token_ids_visible: self.stop_token_ids_visible,
823            stop_token_ids_hidden: self.stop_token_ids_hidden,
824            stop_sequences_visible: self.stop_sequences_visible,
825            stop_sequences_hidden: self.stop_sequences_hidden,
826            stopped: false,
827            state: String::new(),
828        })
829    }
830}