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