dynamo_tokenizers/
fastokens.rs

1// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Fastokens backend using the `fastokens` crate for high-performance BPE encoding.
5//!
6//! This module preserves the existing hybrid behavior: `fastokens` handles encoding and
7//! `HuggingFaceTokenizer` handles decoding. Both are loaded from the same `tokenizer.json`
8//! file.
9
10use std::path::Path;
11
12use rayon::prelude::*;
13
14use super::{
15    EncodeSegment, Encoding, Error, Result, TokenIdType,
16    hf::HuggingFaceTokenizer,
17    traits::{DecodeResult, Decoder, Encoder, Tokenizer},
18};
19
20/// Hybrid tokenizer: fast BPE encoding via `fastokens`, decoding via HuggingFace.
21///
22/// Both backends are loaded from the same `tokenizer.json` file.
23pub struct FastTokenizer {
24    fast_encoder: fastokens::Tokenizer,
25    hf_decoder: HuggingFaceTokenizer,
26}
27
28impl FastTokenizer {
29    pub fn from_file(path: &str) -> Result<Self> {
30        let fast_encoder = fastokens::Tokenizer::from_file(Path::new(path))
31            .map_err(|e| Error::msg(format!("Error loading fastokens tokenizer: {e}")))?;
32        let hf_decoder = HuggingFaceTokenizer::from_file(path)?;
33        Ok(Self {
34            fast_encoder,
35            hf_decoder,
36        })
37    }
38}
39
40impl Encoder for FastTokenizer {
41    fn encode(&self, input: &str) -> Result<Encoding> {
42        let ids = self
43            .fast_encoder
44            .encode(input)
45            .map_err(|e| Error::msg(format!("Fastokens encode error: {e}")))?;
46        Ok(Encoding::Sp(ids))
47    }
48
49    fn encode_batch(&self, inputs: &[&str]) -> Result<Vec<Encoding>> {
50        inputs.par_iter().map(|input| self.encode(input)).collect()
51    }
52
53    fn encode_segments(&self, segments: &[EncodeSegment<'_>]) -> Result<Encoding> {
54        let segments: Vec<fastokens::EncodeSegment<'_>> = segments
55            .iter()
56            .map(|segment| fastokens::EncodeSegment {
57                text: segment.text,
58                allow_special: segment.allow_special,
59            })
60            .collect();
61        let ids = self
62            .fast_encoder
63            .encode_segments(&segments)
64            .map_err(|e| Error::msg(format!("Fastokens segmented encode error: {e}")))?;
65        Ok(Encoding::Sp(ids))
66    }
67}
68
69impl Decoder for FastTokenizer {
70    fn decode(&self, token_ids: &[TokenIdType], skip_special_tokens: bool) -> Result<DecodeResult> {
71        self.hf_decoder.decode(token_ids, skip_special_tokens)
72    }
73}
74
75impl Tokenizer for FastTokenizer {
76    fn validate_prefix_cache(&self) -> Result<()> {
77        Ok(())
78    }
79}
80
81#[cfg(test)]
82mod tests {
83    use super::*;
84    use crate::HuggingFaceTokenizer;
85
86    // Minimal synthetic BPE tokenizer with no normalizer or post-processor --
87    // compatible with fastokens. Vocab covers: H,T,a,d,e,h,i,l,o,r,s,t,w + punctuation.
88    const TOKENIZER_PATH: &str = concat!(
89        env!("CARGO_MANIFEST_DIR"),
90        "/tests/data/minimal-bpe/tokenizer.json"
91    );
92    const SEGMENTED_TOKENIZER_PATH: &str = concat!(
93        env!("CARGO_MANIFEST_DIR"),
94        "/tests/data/sample-models/TinyLlama_v1.1/tokenizer.json"
95    );
96
97    #[test]
98    fn test_fast_encode_decode_roundtrip() {
99        let tokenizer = FastTokenizer::from_file(TOKENIZER_PATH).unwrap();
100        // Encode then decode: verifies both paths execute without error.
101        // With a null decoder, HF inserts spaces between tokens so exact equality
102        // is not expected here -- we just verify the operations succeed and produce
103        // non-empty results.
104        let text = "Hello, world!";
105        let encoding = tokenizer.encode(text).unwrap();
106        assert!(!encoding.token_ids().is_empty());
107        let decoded: String = tokenizer.decode(encoding.token_ids(), true).unwrap().into();
108        assert!(!decoded.is_empty());
109        // The decoded text should contain the same non-space characters
110        let enc_chars: String = text.chars().filter(|c| !c.is_whitespace()).collect();
111        let dec_chars: String = decoded.chars().filter(|c| !c.is_whitespace()).collect();
112        assert_eq!(
113            enc_chars, dec_chars,
114            "non-space characters must be preserved"
115        );
116    }
117
118    #[test]
119    fn test_fast_matches_hf_encoding() {
120        let fast = FastTokenizer::from_file(TOKENIZER_PATH).unwrap();
121        let hf = HuggingFaceTokenizer::from_file(TOKENIZER_PATH).unwrap();
122
123        for text in &["Hello, world!", "Hello", " world", "He llo"] {
124            let fast_ids = fast.encode(text).unwrap();
125            let hf_ids = hf.encode(text).unwrap();
126            assert_eq!(
127                fast_ids.token_ids(),
128                hf_ids.token_ids(),
129                "fastokens and HuggingFace must produce identical token IDs for '{text}'"
130            );
131        }
132    }
133
134    #[test]
135    fn test_fast_batch_encode() {
136        let tokenizer = FastTokenizer::from_file(TOKENIZER_PATH).unwrap();
137        let inputs = &["Hello", " world", "Hello, world!"];
138        let encodings = tokenizer.encode_batch(inputs).unwrap();
139        assert_eq!(encodings.len(), inputs.len());
140        for (enc, input) in encodings.iter().zip(inputs.iter()) {
141            assert!(
142                !enc.token_ids().is_empty(),
143                "encoding for '{input}' must be non-empty"
144            );
145        }
146    }
147
148    #[test]
149    fn test_fast_segmented_encoding_preserves_trust_boundaries() {
150        let tokenizer = FastTokenizer::from_file(SEGMENTED_TOKENIZER_PATH).unwrap();
151        let upstream =
152            fastokens::Tokenizer::from_file(std::path::Path::new(SEGMENTED_TOKENIZER_PATH))
153                .unwrap();
154        let marker = "<s>";
155
156        let trusted = tokenizer
157            .encode_segments(&[EncodeSegment::control(marker)])
158            .unwrap();
159        assert_eq!(
160            trusted.token_ids(),
161            &[upstream.token_to_id(marker).unwrap()],
162            "trusted renderer output must recognize the control token"
163        );
164
165        let ordinary = tokenizer
166            .encode_segments(&[EncodeSegment::ordinary(marker)])
167            .unwrap();
168        assert_ne!(
169            ordinary.token_ids(),
170            trusted.token_ids(),
171            "untrusted content must encode the control-token spelling as ordinary text"
172        );
173
174        let segments = [
175            EncodeSegment::ordinary("hello "),
176            EncodeSegment::control(marker),
177            EncodeSegment::ordinary(marker),
178        ];
179        let upstream_segments = [
180            fastokens::EncodeSegment::ordinary("hello "),
181            fastokens::EncodeSegment::special(marker),
182            fastokens::EncodeSegment::ordinary(marker),
183        ];
184        let actual = tokenizer.encode_segments(&segments).unwrap();
185        let expected = upstream.encode_segments(&upstream_segments).unwrap();
186        assert_eq!(actual.token_ids(), expected);
187
188        assert!(
189            tokenizer
190                .encode_segments(&[])
191                .unwrap()
192                .token_ids()
193                .is_empty()
194        );
195    }
196
197    #[test]
198    fn test_fast_with_decode_stream() {
199        use crate::Tokenizer as TokenizerWrapper;
200        use std::sync::Arc;
201
202        let tokenizer = Arc::new(FastTokenizer::from_file(TOKENIZER_PATH).unwrap());
203        let wrapper = TokenizerWrapper::from(tokenizer);
204
205        // Encode a prompt and a continuation, then step through the decode stream
206        let prompt_ids = wrapper.encode("Hello").unwrap().token_ids().to_vec();
207        let continuation = ", world!";
208        let cont_ids = wrapper.encode(continuation).unwrap().token_ids().to_vec();
209
210        let mut stream = wrapper.decode_stream(&prompt_ids, true);
211        // Accumulate incremental chunks from decode_stream
212        let mut accumulated = String::new();
213        for id in &cont_ids {
214            if let Some(chunk) = stream.step(*id).unwrap() {
215                accumulated.push_str(&chunk);
216            }
217        }
218
219        // DecodeStream uses prompt tokens as context, so the expected text is
220        // decode(prompt + continuation) minus decode(prompt) -- not a bare
221        // decode(continuation) which lacks the surrounding context.
222        let mut all_ids = prompt_ids.clone();
223        all_ids.extend_from_slice(&cont_ids);
224        let full_text: String = wrapper.decode(&all_ids, true).unwrap().into();
225        let prompt_text: String = wrapper.decode(&prompt_ids, true).unwrap().into();
226        let expected = &full_text[prompt_text.len()..];
227        assert_eq!(
228            accumulated, expected,
229            "streamed chunks must equal context-aware decoded continuation"
230        );
231    }
232}
dynamo_tokenizers/fastokens.rs

dynamo_tokenizers/
fastokens.rs
