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