ferrum_interfaces/
tokenizer.rs

1//! Tokenizer interface for text encoding/decoding
2//!
3//! This module provides tokenizer abstractions that are completely separate
4//! from model implementations, supporting incremental decoding and various
5//! tokenization strategies.
6
7use async_trait::async_trait;
8use ferrum_types::{Result, SpecialTokens, TokenId};
9use serde::{Deserialize, Serialize};
10use std::collections::HashMap;
11
12/// Core tokenizer trait for encoding/decoding operations
13pub trait Tokenizer: Send + Sync {
14    /// Encode text to token IDs
15    fn encode(&self, text: &str, add_special: bool) -> Result<Vec<TokenId>>;
16
17    /// Decode token IDs to text
18    fn decode(&self, tokens: &[TokenId], skip_special: bool) -> Result<String>;
19
20    /// Incremental decode: given previous tokens and new token, return only the new text
21    /// This is crucial for streaming applications to avoid re-decoding all tokens
22    fn decode_incremental(&self, prev: &[TokenId], next: TokenId) -> Result<String>;
23
24    /// Get vocabulary size
25    fn vocab_size(&self) -> usize;
26
27    /// Get special tokens configuration  
28    fn special_tokens(&self) -> &SpecialTokens;
29
30    /// Get token ID for a specific text (if exists in vocabulary)
31    fn token_id(&self, text: &str) -> Option<TokenId>;
32
33    /// Get text for a specific token ID
34    fn token_text(&self, token_id: TokenId) -> Option<&str>;
35
36    /// Return the context-free byte surface represented by one vocabulary token.
37    ///
38    /// This is distinct from decoding a one-token sequence: byte-level BPE
39    /// vocabularies may split one UTF-8 scalar across multiple tokens, so a
40    /// string decoder must replace an incomplete fragment. Grammar/token-trie
41    /// consumers need the original bytes instead. Tokenizers with a byte-level
42    /// vocabulary should override this method.
43    fn token_bytes(&self, token_id: TokenId) -> Option<Vec<u8>> {
44        self.decode(&[token_id], false)
45            .ok()
46            .map(String::into_bytes)
47            .or_else(|| {
48                self.token_text(token_id)
49                    .map(|text| text.as_bytes().to_vec())
50            })
51    }
52
53    /// Check if token is a special token
54    fn is_special_token(&self, token_id: TokenId) -> bool {
55        let special = self.special_tokens();
56        let fallback = TokenId::MAX;
57        token_id == special.bos_token.unwrap_or(fallback)
58            || token_id == special.eos_token.unwrap_or(fallback)
59            || token_id == special.unk_token.unwrap_or(fallback)
60            || token_id == special.pad_token.unwrap_or(fallback)
61    }
62
63    /// Apply chat template if supported
64    fn apply_chat_template(&self, messages: &[ChatMessage]) -> Result<String> {
65        // Default implementation: just concatenate messages
66        let mut result = String::new();
67        for msg in messages {
68            result.push_str(&format!("{}: {}\n", msg.role, msg.content));
69        }
70        Ok(result.trim_end().to_string())
71    }
72
73    /// Get tokenizer information
74    fn info(&self) -> TokenizerInfo;
75}
76
77/// Asynchronous tokenizer operations for I/O-bound tokenization
78#[async_trait]
79pub trait AsyncTokenizer: Tokenizer {
80    /// Asynchronous encoding (useful for very large texts)
81    async fn encode_async(&self, text: &str, add_special: bool) -> Result<Vec<TokenId>>;
82
83    /// Asynchronous decoding
84    async fn decode_async(&self, tokens: &[TokenId], skip_special: bool) -> Result<String>;
85
86    /// Batch encoding for multiple texts
87    async fn encode_batch(&self, texts: &[&str], add_special: bool) -> Result<Vec<Vec<TokenId>>>;
88
89    /// Batch decoding for multiple token sequences
90    async fn decode_batch(
91        &self,
92        token_sequences: &[&[TokenId]],
93        skip_special: bool,
94    ) -> Result<Vec<String>>;
95}
96
97/// Advanced tokenizer capabilities
98pub trait TokenizerCapabilities: Tokenizer {
99    /// Get token probability/likelihood for text
100    fn token_probability(&self, text: &str, token_id: TokenId) -> Option<f32>;
101
102    /// Get all possible tokens for a prefix
103    fn get_prefix_tokens(&self, prefix: &str) -> Result<Vec<TokenId>>;
104
105    /// Check if sequence can be extended with token
106    fn can_extend(&self, tokens: &[TokenId], next_token: TokenId) -> bool;
107
108    /// Get token type (word, subword, punctuation, etc.)
109    fn token_type(&self, token_id: TokenId) -> TokenType;
110
111    /// Normalize text before tokenization
112    fn normalize_text(&self, text: &str) -> String;
113
114    /// Pre-tokenize text (split into words/subwords)
115    fn pre_tokenize(&self, text: &str) -> Vec<String>;
116}
117
118/// Tokenizer factory for creating tokenizer instances
119#[async_trait]
120pub trait TokenizerFactory: Send + Sync {
121    /// Load tokenizer from file path
122    async fn load_from_file(&self, path: &str) -> Result<Box<dyn Tokenizer>>;
123
124    /// Load tokenizer from bytes
125    async fn load_from_bytes(&self, data: &[u8]) -> Result<Box<dyn Tokenizer>>;
126
127    /// Load tokenizer from Hugging Face Hub
128    async fn load_from_hub(
129        &self,
130        repo_id: &str,
131        revision: Option<&str>,
132    ) -> Result<Box<dyn Tokenizer>>;
133
134    /// Create tokenizer from configuration
135    async fn create_from_config(&self, config: &TokenizerConfig) -> Result<Box<dyn Tokenizer>>;
136
137    /// Get supported tokenizer types
138    fn supported_types(&self) -> Vec<TokenizerType>;
139}
140
141/// Tokenizer information and metadata
142#[derive(Debug, Clone, Serialize, Deserialize)]
143pub struct TokenizerInfo {
144    /// Tokenizer type/algorithm
145    pub tokenizer_type: TokenizerType,
146    /// Vocabulary size
147    pub vocab_size: usize,
148    /// Special tokens
149    pub special_tokens: SpecialTokens,
150    /// Whether tokenizer supports incremental decoding efficiently
151    pub supports_incremental: bool,
152    /// Whether tokenizer supports chat templates
153    pub supports_chat_template: bool,
154    /// Maximum token length
155    pub max_token_length: Option<usize>,
156    /// Model name or identifier this tokenizer was trained for
157    pub model_name: Option<String>,
158}
159
160/// Tokenizer types/algorithms
161#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
162pub enum TokenizerType {
163    /// Byte-Pair Encoding
164    BPE,
165    /// WordPiece (BERT-style)
166    WordPiece,
167    /// SentencePiece
168    SentencePiece,
169    /// Tiktoken (GPT family)
170    Tiktoken,
171    /// Custom implementation
172    Custom,
173}
174
175/// Token types for classification
176#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
177pub enum TokenType {
178    /// Regular word token
179    Word,
180    /// Subword token  
181    Subword,
182    /// Punctuation token
183    Punctuation,
184    /// Number token
185    Number,
186    /// Special/control token
187    Special,
188    /// Unknown token
189    Unknown,
190}
191
192/// Chat message for template application
193#[derive(Debug, Clone, Serialize, Deserialize)]
194pub struct ChatMessage {
195    /// Message role (user, assistant, system, etc.)
196    pub role: String,
197    /// Message content
198    pub content: String,
199    /// Additional metadata
200    pub metadata: HashMap<String, serde_json::Value>,
201}
202
203impl ChatMessage {
204    /// Create user message
205    pub fn user(content: impl Into<String>) -> Self {
206        Self {
207            role: "user".to_string(),
208            content: content.into(),
209            metadata: HashMap::new(),
210        }
211    }
212
213    /// Create assistant message
214    pub fn assistant(content: impl Into<String>) -> Self {
215        Self {
216            role: "assistant".to_string(),
217            content: content.into(),
218            metadata: HashMap::new(),
219        }
220    }
221
222    /// Create system message
223    pub fn system(content: impl Into<String>) -> Self {
224        Self {
225            role: "system".to_string(),
226            content: content.into(),
227            metadata: HashMap::new(),
228        }
229    }
230}
231
232/// Tokenizer configuration
233#[derive(Debug, Clone, Serialize, Deserialize)]
234pub struct TokenizerConfig {
235    /// Tokenizer type
236    pub tokenizer_type: TokenizerType,
237    /// Path to tokenizer files
238    pub path: String,
239    /// Whether to add special tokens during encoding
240    pub add_special_tokens: bool,
241    /// Whether to use fast tokenization (if available)
242    pub use_fast: bool,
243    /// Truncation configuration
244    pub truncation: Option<TruncationConfig>,
245    /// Padding configuration
246    pub padding: Option<PaddingConfig>,
247    /// Chat template (if any)
248    pub chat_template: Option<String>,
249    /// Additional tokenizer-specific options
250    pub extra_options: HashMap<String, serde_json::Value>,
251}
252
253/// Truncation configuration
254#[derive(Debug, Clone, Serialize, Deserialize)]
255pub struct TruncationConfig {
256    /// Maximum sequence length
257    pub max_length: usize,
258    /// Truncation strategy
259    pub strategy: TruncationStrategy,
260    /// Stride for sliding window truncation
261    pub stride: Option<usize>,
262}
263
264/// Truncation strategies
265#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
266pub enum TruncationStrategy {
267    /// Remove tokens from the end
268    TruncateEnd,
269    /// Remove tokens from the beginning  
270    TruncateStart,
271    /// Remove tokens from both ends equally
272    TruncateBoth,
273    /// Sliding window approach
274    SlidingWindow,
275}
276
277/// Padding configuration
278#[derive(Debug, Clone, Serialize, Deserialize)]
279pub struct PaddingConfig {
280    /// Padding strategy
281    pub strategy: PaddingStrategy,
282    /// Padding token ID
283    pub token_id: TokenId,
284    /// Target length (if fixed padding)
285    pub length: Option<usize>,
286    /// Padding direction
287    pub direction: PaddingDirection,
288}
289
290/// Padding strategies
291#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
292pub enum PaddingStrategy {
293    /// No padding
294    None,
295    /// Pad to longest sequence in batch
296    Longest,
297    /// Pad to multiple of specified value
298    MultipleOf(usize),
299    /// Pad to fixed length
300    Fixed,
301}
302
303/// Padding direction
304#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
305pub enum PaddingDirection {
306    /// Pad on the right
307    Right,
308    /// Pad on the left
309    Left,
310}
311
312/// Incremental tokenizer state for streaming
313pub trait IncrementalTokenizer: Tokenizer {
314    /// Tokenizer state for incremental operations
315    type State: Send + Sync;
316
317    /// Create initial state for incremental decoding
318    fn create_state(&self) -> Self::State;
319
320    /// Add token to state and get incremental text
321    fn decode_incremental_with_state(
322        &self,
323        state: &mut Self::State,
324        token: TokenId,
325    ) -> Result<String>;
326
327    /// Reset state to initial condition
328    fn reset_state(&self, state: &mut Self::State);
329
330    /// Get all decoded text from current state
331    fn get_decoded_text(&self, state: &Self::State) -> String;
332}
333
334/// Text processing utilities
335pub trait TextProcessor: Send + Sync {
336    /// Clean and normalize text for tokenization
337    fn preprocess(&self, text: &str) -> String;
338
339    /// Post-process decoded text
340    fn postprocess(&self, text: &str) -> String;
341
342    /// Detect language of text (if supported)
343    fn detect_language(&self, text: &str) -> Option<String>;
344
345    /// Split text into sentences
346    fn sentence_split(&self, text: &str) -> Vec<String>;
347
348    /// Count approximate tokens without full tokenization
349    fn estimate_token_count(&self, text: &str) -> usize;
350}
351
352/// Tokenizer performance statistics
353#[derive(Debug, Clone, Serialize, Deserialize)]
354pub struct TokenizerStats {
355    /// Total encoding operations
356    pub encode_operations: u64,
357    /// Total decoding operations  
358    pub decode_operations: u64,
359    /// Total tokens processed
360    pub tokens_processed: u64,
361    /// Average encoding time per character (microseconds)
362    pub avg_encode_time_per_char_us: f64,
363    /// Average decoding time per token (microseconds)
364    pub avg_decode_time_per_token_us: f64,
365    /// Cache hit rate for incremental decoding
366    pub incremental_cache_hit_rate: f32,
367}
368
369/// Tokenizer registry for managing multiple tokenizers
370pub trait TokenizerRegistry: Send + Sync {
371    /// Register a tokenizer with a name
372    fn register(&mut self, name: &str, tokenizer: Box<dyn Tokenizer>) -> Result<()>;
373
374    /// Get tokenizer by name
375    fn get(&self, name: &str) -> Option<&dyn Tokenizer>;
376
377    /// Remove tokenizer by name
378    fn remove(&mut self, name: &str) -> Option<Box<dyn Tokenizer>>;
379
380    /// List all registered tokenizer names
381    fn list_names(&self) -> Vec<String>;
382
383    /// Check if tokenizer exists
384    fn contains(&self, name: &str) -> bool;
385}
ferrum_interfaces/tokenizer.rs

ferrum_interfaces/
tokenizer.rs
