deepseek_recipe_encoding/
tokenizer.rs

1//! Encode rendered prompts into token IDs.
2
3use std::collections::HashMap;
4use std::ops::Range;
5use std::sync::Arc;
6
7use tokenizers::{AddedVocabulary, Tokenizer};
8
9/// An encoder that maps prompt text to model token IDs.
10///
11/// Implemented for [`PromptTokenizer`]; other tokenizers can implement it to
12/// provide token encoding for a prompt encoding.
13pub trait TokenizerEncoder: Send + Sync {
14    /// Encode `text` into token IDs.
15    ///
16    /// `special_ranges` are the byte ranges of `text` that hold explicit
17    /// special tokens. The caller must pass ranges in ascending order that do
18    /// not overlap, lie within `text`, and start and end on `char` boundaries;
19    /// [`RenderedPrompt::special_ranges`] meets these conditions.
20    ///
21    /// The implementation must encode each range as the token ID of its
22    /// spelling and the plain text between the ranges without the added
23    /// vocabulary, so a special-token spelling in plain text encodes as the
24    /// token IDs of its characters. The implementation must add no tokens,
25    /// such as a BOS token, beyond those that the text encodes to.
26    ///
27    /// # Errors
28    ///
29    /// Returns an error message when the implementation cannot encode `text`.
30    ///
31    /// [`RenderedPrompt::special_ranges`]: crate::RenderedPrompt::special_ranges
32    fn encode_ids(&self, text: &str, special_ranges: &[Range<usize>]) -> Result<Vec<u32>, String>;
33}
34
35/// A tokenizer prepared for prompt encoding.
36///
37/// It holds a copy of the tokenizer without its added vocabulary, which
38/// encodes plain text, and the added vocabulary, which supplies special token
39/// IDs.
40#[derive(Clone)]
41pub struct PromptTokenizer {
42    plain_tokenizer: Tokenizer,
43    added_vocabulary: HashMap<String, u32>,
44}
45
46impl PromptTokenizer {
47    /// Prepare `tokenizer` for prompt encoding. This clones the tokenizer once.
48    pub fn new(tokenizer: &Tokenizer) -> Self {
49        let added_vocabulary = tokenizer
50            .get_added_vocabulary()
51            .get_vocab()
52            .iter()
53            .map(|(token, &id)| (token.clone(), id))
54            .collect();
55        let mut plain_tokenizer = tokenizer.clone();
56        plain_tokenizer.with_added_vocabulary(AddedVocabulary::new());
57        Self {
58            plain_tokenizer,
59            added_vocabulary,
60        }
61    }
62
63    fn special_token_id(&self, token: &str) -> Result<u32, String> {
64        self.added_vocabulary
65            .get(token)
66            .copied()
67            .ok_or_else(|| format!("special token {token:?} is not in the added vocabulary"))
68    }
69}
70
71impl TokenizerEncoder for PromptTokenizer {
72    /// Encode the plain text regions as one batch with the tokenizer without
73    /// its added vocabulary, and each special range as the added-vocabulary
74    /// token ID of its spelling. The batch runs in parallel unless the
75    /// `TOKENIZERS_PARALLELISM` environment variable disables it. The
76    /// tokenizer adds no tokens, such as a BOS token, beyond those that the
77    /// text encodes to, provided that it has no padding configured.
78    ///
79    /// # Errors
80    ///
81    /// Returns an error message when the tokenizer fails to encode a plain
82    /// text region or the spelling of a special range is not in the added
83    /// vocabulary.
84    ///
85    /// # Panics
86    ///
87    /// Panics if a range starts before the end of the previous range, ends
88    /// before it starts, extends past the end of `text`, or does not start and
89    /// end on `char` boundaries.
90    fn encode_ids(&self, text: &str, special_ranges: &[Range<usize>]) -> Result<Vec<u32>, String> {
91        // The plain region before each special range, then the one after the
92        // last range; a region can be empty.
93        let mut regions = Vec::with_capacity(special_ranges.len() + 1);
94        let mut start = 0;
95        for range in special_ranges {
96            regions.push(&text[start..range.start]);
97            start = range.end;
98        }
99        regions.push(&text[start..]);
100        let encodings = self
101            .plain_tokenizer
102            .encode_batch_fast(regions, false)
103            .map_err(|error| error.to_string())?;
104
105        let mut ids = Vec::new();
106        let mut encodings = encodings.iter();
107        for range in special_ranges {
108            if let Some(encoding) = encodings.next() {
109                ids.extend_from_slice(encoding.get_ids());
110            }
111            ids.push(self.special_token_id(&text[range.clone()])?);
112        }
113        if let Some(encoding) = encodings.next() {
114            ids.extend_from_slice(encoding.get_ids());
115        }
116        Ok(ids)
117    }
118}
119
120impl<T: TokenizerEncoder + ?Sized> TokenizerEncoder for Arc<T> {
121    fn encode_ids(&self, text: &str, special_ranges: &[Range<usize>]) -> Result<Vec<u32>, String> {
122        (**self).encode_ids(text, special_ranges)
123    }
124}
deepseek_recipe_encoding/tokenizer.rs

deepseek_recipe_encoding/
tokenizer.rs
