[smg] remove dead tokenizer code (#17722)
This commit is contained in:
Simo Lin
@@ -392,55 +392,3 @@ pub async fn get_tokenizer_status(context: &Arc<AppContext>, tokenizer_id: &str)
|
||||
"not_found",
|
||||
)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::tokenizer::mock::MockTokenizer;
|
||||
|
||||
fn create_test_registry() -> Arc<TokenizerRegistry> {
|
||||
let registry = Arc::new(TokenizerRegistry::new());
|
||||
let id = TokenizerRegistry::generate_id();
|
||||
registry.register(
|
||||
&id,
|
||||
"test-model",
|
||||
"test-source",
|
||||
Arc::new(MockTokenizer::new()),
|
||||
);
|
||||
registry
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_get_tokenizer_exact_match() {
|
||||
let registry = create_test_registry();
|
||||
let result = get_tokenizer(®istry, "test-model");
|
||||
assert!(result.is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_get_tokenizer_unknown_model_fallback() {
|
||||
let registry = create_test_registry();
|
||||
let result = get_tokenizer(®istry, UNKNOWN_MODEL_ID);
|
||||
assert!(result.is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_get_tokenizer_not_found() {
|
||||
let registry = create_test_registry();
|
||||
let result = get_tokenizer(®istry, "nonexistent");
|
||||
match result {
|
||||
Err(e) => assert!(e.contains("not found")),
|
||||
Ok(_) => panic!("Expected error"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_get_tokenizer_empty_registry() {
|
||||
let registry = Arc::new(TokenizerRegistry::new());
|
||||
let result = get_tokenizer(®istry, "any");
|
||||
match result {
|
||||
Err(e) => assert!(e.contains("No tokenizers available")),
|
||||
Ok(_) => panic!("Expected error"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,197 +0,0 @@
|
||||
# Tokenizer Module
|
||||
|
||||
## Overview
|
||||
The `sgl-model-gateway` tokenizer subsystem exposes a single `Tokenizer` facade around multiple backends
|
||||
(Hugging Face JSON tokenizers, OpenAI/tiktoken models, and an in-memory mock). It packages the
|
||||
shared behaviours needed by the router–encoding user text, incrementally decoding streamed tokens,
|
||||
tracking per-request state, and detecting stop conditions—behind trait objects so the rest of the
|
||||
router can remain backend-agnostic.
|
||||
|
||||
Key capabilities:
|
||||
- trait-based split between `Encoder`, `Decoder`, and `Tokenizer` for shared APIs across backends
|
||||
- Hugging Face tokenizer loading (with optional chat templates) and HF Hub downloads
|
||||
- heuristic selection of OpenAI/tiktoken encodings for GPT model names
|
||||
- incremental decoding utilities (`DecodeStream`, `Sequence`) that handle UTF-8 boundaries
|
||||
- stop sequence handling via `StopSequenceDecoder` with token-level and string-level triggers
|
||||
- optional Jinja2 chat-template rendering that matches Hugging Face semantics
|
||||
|
||||
The implementation deliberately keeps the surface area small—metrics, batching, or SentencePiece
|
||||
support mentioned in earlier drafts do **not** exist today. This document reflects the actual code
|
||||
as of `sgl-model-gateway/src/tokenizer/*`.
|
||||
|
||||
## Source Map
|
||||
- `mod.rs` – module exports and the `Tokenizer` wrapper around `Arc<dyn Tokenizer>`
|
||||
- `traits.rs` – shared traits and the `Encoding`/`SpecialTokens` helper types
|
||||
- `factory.rs` – backend discovery, file/model heuristics, and tokio-aware creation helpers
|
||||
- `hub.rs` – Hugging Face Hub downloads via `hf_hub`
|
||||
- `huggingface.rs` – wrapper over `tokenizers::Tokenizer`, chat template loading, vocab access
|
||||
- `tiktoken.rs` – wrapper over `tiktoken-rs` encoders for OpenAI model families
|
||||
- `chat_template.rs` – AST-driven Jinja template inspection and rendering utilities
|
||||
- `sequence.rs` – stateful incremental decoding helper used by router sequences
|
||||
- `stream.rs` – stateless streaming decoder that yields textual chunks from token streams
|
||||
- `stop.rs` – stop-sequence detection with "jail" buffering and a builder API
|
||||
- `mock.rs` – lightweight tokenizer used by unit tests
|
||||
- `tests.rs` – smoke tests covering the trait facade and helpers (largely with the mock backend)
|
||||
|
||||
## Core Traits and Types (`traits.rs`)
|
||||
- `Encoder`, `Decoder`, and `Tokenizer` traits stay `Send + Sync` so instances can be shared across
|
||||
threads. Concrete backends implement the minimal methods: `encode`, `encode_batch`, `decode`,
|
||||
`vocab_size`, special-token lookup, and optional token↔id conversions.
|
||||
- `Encoding` wraps backend-specific results: `Hf` holds the Hugging Face encoding object,
|
||||
`Sp` is a plain ID vector reserved for future SentencePiece support, and `Tiktoken` stores u32 IDs
|
||||
from `tiktoken-rs`. `Encoding::token_ids()` is the zero-copy accessor used everywhere.
|
||||
- `SpecialTokens` collects optional BOS/EOS/etc. markers so upstream code can make backend-agnostic
|
||||
decisions.
|
||||
- `Tokenizer` (in `mod.rs`) is a thin `Arc<dyn Tokenizer>` newtype that exposes convenience methods
|
||||
(`encode`, `decode`, `decode_stream`, etc.) while keeping cloning cheap.
|
||||
|
||||
## Backend Implementations
|
||||
### HuggingFaceTokenizer (`huggingface.rs`)
|
||||
- Loads `tokenizer.json` (or similar) using `tokenizers::Tokenizer::from_file`.
|
||||
- Caches vocab forward and reverse maps for `token_to_id`/`id_to_token` support.
|
||||
- Extracts special tokens using common patterns (e.g. `<s>`, `[CLS]`).
|
||||
- Supports optional chat templates: either auto-discovered next to the tokenizer via
|
||||
`tokenizer_config.json` or overridable with an explicit template path.
|
||||
- Exposes `apply_chat_template` which renders a minijinja template given JSON message payloads and
|
||||
template parameters.
|
||||
|
||||
### TiktokenTokenizer (`tiktoken.rs`)
|
||||
- Wraps the `tiktoken-rs` `CoreBPE` builders (`cl100k_base`, `p50k_base`, `p50k_edit`, `r50k_base`).
|
||||
- `from_model_name` heuristically maps OpenAI model IDs (e.g. `gpt-4`, `text-davinci-003`) to those
|
||||
bases. Unknown model names return an error rather than silently defaulting.
|
||||
- Implements encode/decode operations; batch encode simply iterates sequentially.
|
||||
- Provides approximate vocab sizes and common GPT special tokens. Direct token↔id lookup is not
|
||||
implemented—the underlying library does not expose that mapping.
|
||||
|
||||
### MockTokenizer (`mock.rs`)
|
||||
- Purely for tests; hard-codes a tiny vocabulary and simple whitespace tokenization.
|
||||
- Implements the same trait surface so helpers can be exercised without pulling real tokenizer data.
|
||||
|
||||
## Factory and Backend Discovery (`factory.rs`)
|
||||
- `create_tokenizer{,_async}` accept either a filesystem path or a model identifier. Logic:
|
||||
1. Paths are loaded directly; the file extension (or JSON autodetection) selects the backend.
|
||||
2. Strings that look like OpenAI model names (`gpt-*`, `davinci`, `curie`, `babbage`, `ada`) use
|
||||
`TiktokenTokenizer`.
|
||||
3. Everything else attempts a Hugging Face Hub download via `download_tokenizer_from_hf`.
|
||||
- Chat templates can be injected with `create_tokenizer_with_chat_template`.
|
||||
- Async creation uses `tokio` for network access. The blocking variant reuses or spins up a runtime
|
||||
when called from synchronous contexts.
|
||||
- SentencePiece (`.model`) and GGUF files are detected but currently return a clear `not supported`
|
||||
error.
|
||||
|
||||
## Hugging Face Hub Integration (`hub.rs`)
|
||||
- Uses the async `hf_hub` API to list and download tokenizer-related files
|
||||
(`tokenizer.json`, `merges.txt`, `.model`, etc.), filtering out weights and docs.
|
||||
- The helper returns the HF cache directory containing the fetched files; the factory then loads
|
||||
from disk using standard file paths.
|
||||
- Honour the `HF_TOKEN` environment variable for private or rate-limited models. Without it the
|
||||
download may fail with an authorization error.
|
||||
|
||||
## Chat Template Support (`chat_template.rs`)
|
||||
- Detects whether a template expects raw string content or the structured OpenAI-style `content`
|
||||
list by walking the minijinja AST. This matches the Python-side detection logic used elsewhere in
|
||||
SGLang.
|
||||
- `ChatTemplateProcessor` (constructed per call) renders templates against JSON `messages` and
|
||||
`ChatTemplateParams` (system prompt, tools, EOS token handling, etc.). Errors surface as
|
||||
`anyhow::Error`, keeping parity with Hugging Face error messages.
|
||||
- The tokenizer wrapper stores both the template string and its detected content format so callers
|
||||
can pre-transform message content correctly.
|
||||
|
||||
## Streaming and Stateful Helpers
|
||||
### `DecodeStream` (`stream.rs`)
|
||||
- Maintains a sliding window (`prefix_offset`, `read_offset`) over accumulated token IDs.
|
||||
- Each `step` decodes the known prefix and the new slice; when the new slice produces additional
|
||||
UTF-8 text (and does not end in the replacement character `<60>`), it returns the incremental chunk
|
||||
and updates offsets. Otherwise it returns `None` and waits for more tokens.
|
||||
- `step_batch` and `flush` offer convenience for batching and draining remaining text.
|
||||
|
||||
### `Sequence` (`sequence.rs`)
|
||||
- Holds per-request decoding state: accumulated IDs plus offsets mirroring `DecodeStream`.
|
||||
- `append_text` encodes extra prompt text; `append_token` decodes incremental output while
|
||||
respecting UTF-8 boundaries and replacing stray `<60>` characters.
|
||||
- Designed for integration with router sequence management where decoded text must be replayed.
|
||||
|
||||
### `StopSequenceDecoder` (`stop.rs`)
|
||||
- Extends the incremental decoding approach with a "jail" buffer that holds potential partial
|
||||
matches against configured stop sequences.
|
||||
- Supports both token-level stops (visible or hidden) and arbitrary string sequences. When a string
|
||||
stop is configured, the decoder emits only the safe prefix and keeps a suffix jailed until it can
|
||||
decide whether it completes a stop sequence.
|
||||
- Provides `StopSequenceDecoderBuilder` for ergonomic configuration and exposes `process_token`,
|
||||
`process_tokens`, `flush`, `reset`, and `is_stopped` helpers.
|
||||
|
||||
## Testing
|
||||
- Unit tests cover the mock tokenizer, the `Tokenizer` wrapper, incremental decoding helpers, and
|
||||
stop-sequence behaviour (`tests.rs`, `sequence.rs`, `stop.rs`, `tiktoken.rs`, `factory.rs`,
|
||||
`hub.rs`). Network-dependent Hugging Face downloads are exercised behind a best-effort async test
|
||||
that skips in CI without credentials.
|
||||
- Use `cargo test -p sgl-model-gateway tokenizer` to run the module’s test suite.
|
||||
|
||||
## Known Limitations & Future Work
|
||||
- SentencePiece (`.model`) and GGUF tokenizers are detected but deliberately unimplemented.
|
||||
- `Encoding::Sp` exists for future SentencePiece support but currently behaves as a simple `Vec<u32>`.
|
||||
- `TiktokenTokenizer` cannot map individual tokens/IDs; the underlying library would need to expose
|
||||
its vocabulary to implement `token_to_id`/`id_to_token`.
|
||||
- There is no metrics or batching layer inside this module; the router records metrics elsewhere.
|
||||
- Dynamic batching / sequence pooling code that earlier READMEs mentioned never landed in Rust.
|
||||
|
||||
## Usage Examples
|
||||
```rust
|
||||
use std::sync::Arc;
|
||||
use smg::tokenizer::{
|
||||
create_tokenizer, SequenceDecoderOutput, StopSequenceDecoderBuilder, Tokenizer,
|
||||
};
|
||||
|
||||
// Load a tokenizer from disk (Hugging Face JSON)
|
||||
let tokenizer = Tokenizer::from_file("/path/to/tokenizer.json")?;
|
||||
let encoding = tokenizer.encode("Hello, world!")?;
|
||||
assert!(!encoding.token_ids().is_empty());
|
||||
|
||||
// Auto-detect OpenAI GPT tokenizer
|
||||
let openai = create_tokenizer("gpt-4")?;
|
||||
let text = openai.decode(&[1, 2, 3], true)?;
|
||||
|
||||
// Incremental decoding with stop sequences
|
||||
let mut stream = tokenizer.decode_stream(&[], true);
|
||||
let mut stop = StopSequenceDecoderBuilder::new(Arc::clone(&tokenizer))
|
||||
.stop_sequence("\nHuman:")
|
||||
.build();
|
||||
for &token in encoding.token_ids() {
|
||||
if let Some(chunk) = stream.step(token)? {
|
||||
match stop.process_token(token)? {
|
||||
SequenceDecoderOutput::Text(t) => println!("{}", t),
|
||||
SequenceDecoderOutput::StoppedWithText(t) => {
|
||||
println!("{}", t);
|
||||
break;
|
||||
}
|
||||
SequenceDecoderOutput::Held | SequenceDecoderOutput::Stopped => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
```rust
|
||||
// Apply a chat template when one is bundled with the tokenizer
|
||||
use smg::tokenizer::{chat_template::ChatTemplateParams, HuggingFaceTokenizer};
|
||||
|
||||
let mut hf = HuggingFaceTokenizer::from_file_with_chat_template(
|
||||
"./tokenizer.json",
|
||||
Some("./chat_template.jinja"),
|
||||
)?;
|
||||
let messages = vec![
|
||||
serde_json::json!({"role": "system", "content": "You are concise."}),
|
||||
serde_json::json!({"role": "user", "content": "Summarise Rust traits."}),
|
||||
];
|
||||
let prompt = hf.apply_chat_template(
|
||||
&messages,
|
||||
ChatTemplateParams {
|
||||
add_generation_prompt: true,
|
||||
continue_final_message: false,
|
||||
tools: None,
|
||||
documents: None,
|
||||
template_kwargs: None,
|
||||
},
|
||||
)?;
|
||||
```
|
||||
|
||||
Set `HF_TOKEN` in the environment if you need to download private models from the Hugging Face Hub.
|
||||
106
sgl-model-gateway/src/tokenizer/cache/fingerprint.rs
vendored
106
sgl-model-gateway/src/tokenizer/cache/fingerprint.rs
vendored
@@ -1,106 +0,0 @@
|
||||
//! Tokenizer Fingerprinting for Cache Invalidation
|
||||
//!
|
||||
//! Creates a unique fingerprint of a tokenizer's configuration to detect
|
||||
//! when the tokenizer has changed and the cache needs to be cleared.
|
||||
|
||||
use std::{
|
||||
collections::hash_map::DefaultHasher,
|
||||
hash::{Hash, Hasher},
|
||||
};
|
||||
|
||||
use super::super::traits::Tokenizer;
|
||||
|
||||
/// A fingerprint of a tokenizer's configuration
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
|
||||
pub struct TokenizerFingerprint {
|
||||
/// Size of the vocabulary
|
||||
pub vocab_size: usize,
|
||||
/// Hash of a sample of vocabulary tokens (for speed)
|
||||
pub vocab_hash: u64,
|
||||
/// Hash of special tokens
|
||||
pub special_tokens_hash: u64,
|
||||
}
|
||||
|
||||
impl TokenizerFingerprint {
|
||||
/// Create a fingerprint from a tokenizer
|
||||
pub fn from_tokenizer(tokenizer: &dyn Tokenizer) -> Self {
|
||||
let vocab_size = tokenizer.vocab_size();
|
||||
let vocab_hash = Self::compute_vocab_hash(tokenizer);
|
||||
let special_tokens_hash = Self::compute_special_tokens_hash(tokenizer);
|
||||
|
||||
Self {
|
||||
vocab_size,
|
||||
vocab_hash,
|
||||
special_tokens_hash,
|
||||
}
|
||||
}
|
||||
|
||||
/// Compute a hash of the vocabulary by sampling tokens
|
||||
fn compute_vocab_hash(tokenizer: &dyn Tokenizer) -> u64 {
|
||||
let mut hasher = DefaultHasher::new();
|
||||
let vocab_size = tokenizer.vocab_size();
|
||||
|
||||
// Sample up to 1000 tokens for speed
|
||||
let sample_size = vocab_size.min(1000);
|
||||
let step = if sample_size > 0 {
|
||||
vocab_size / sample_size
|
||||
} else {
|
||||
1
|
||||
};
|
||||
|
||||
for i in (0..vocab_size).step_by(step.max(1)) {
|
||||
if let Some(token) = tokenizer.id_to_token(i as u32) {
|
||||
token.hash(&mut hasher);
|
||||
}
|
||||
}
|
||||
|
||||
hasher.finish()
|
||||
}
|
||||
|
||||
/// Compute a hash of special tokens
|
||||
fn compute_special_tokens_hash(tokenizer: &dyn Tokenizer) -> u64 {
|
||||
let mut hasher = DefaultHasher::new();
|
||||
let special_tokens = tokenizer.get_special_tokens();
|
||||
|
||||
special_tokens.bos_token.hash(&mut hasher);
|
||||
special_tokens.eos_token.hash(&mut hasher);
|
||||
special_tokens.unk_token.hash(&mut hasher);
|
||||
special_tokens.sep_token.hash(&mut hasher);
|
||||
special_tokens.pad_token.hash(&mut hasher);
|
||||
special_tokens.cls_token.hash(&mut hasher);
|
||||
special_tokens.mask_token.hash(&mut hasher);
|
||||
special_tokens.additional_special_tokens.hash(&mut hasher);
|
||||
|
||||
hasher.finish()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::tokenizer::mock::MockTokenizer;
|
||||
|
||||
#[test]
|
||||
fn test_fingerprint_equality() {
|
||||
let tokenizer1 = MockTokenizer::new();
|
||||
let tokenizer2 = MockTokenizer::new();
|
||||
|
||||
let fp1 = TokenizerFingerprint::from_tokenizer(&tokenizer1);
|
||||
let fp2 = TokenizerFingerprint::from_tokenizer(&tokenizer2);
|
||||
|
||||
// Same tokenizer config should produce same fingerprint
|
||||
assert_eq!(fp1, fp2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_fingerprint_consistency() {
|
||||
let tokenizer = MockTokenizer::new();
|
||||
|
||||
let fp1 = TokenizerFingerprint::from_tokenizer(&tokenizer);
|
||||
let fp2 = TokenizerFingerprint::from_tokenizer(&tokenizer);
|
||||
|
||||
// Fingerprint should be consistent
|
||||
assert_eq!(fp1, fp2);
|
||||
assert_eq!(fp1.vocab_size, tokenizer.vocab_size());
|
||||
}
|
||||
}
|
||||
246
sgl-model-gateway/src/tokenizer/cache/l0.rs
vendored
246
sgl-model-gateway/src/tokenizer/cache/l0.rs
vendored
@@ -1,246 +0,0 @@
|
||||
//! L0 Cache: Whole-string exact match cache
|
||||
//!
|
||||
//! This is the simplest and most effective cache layer.
|
||||
//! Key: input string → Value: full encoding result (Arc-wrapped for zero-copy cache hits)
|
||||
//!
|
||||
//! Expected hit rate: 60-90% for workloads with repeated system prompts
|
||||
|
||||
use std::sync::{
|
||||
atomic::{AtomicU64, Ordering},
|
||||
Arc,
|
||||
};
|
||||
|
||||
use dashmap::DashMap;
|
||||
|
||||
use super::super::traits::Encoding;
|
||||
|
||||
/// L0 cache implementation using DashMap for lock-free reads
|
||||
/// Uses Arc<Encoding> internally to provide zero-copy cache hits
|
||||
pub struct L0Cache {
|
||||
/// The cache map: input string → Arc-wrapped encoding for cheap cloning
|
||||
map: Arc<DashMap<String, Arc<Encoding>>>,
|
||||
/// Maximum number of entries before eviction
|
||||
max_entries: usize,
|
||||
/// Cache hit counter
|
||||
hits: AtomicU64,
|
||||
/// Cache miss counter
|
||||
misses: AtomicU64,
|
||||
}
|
||||
|
||||
impl L0Cache {
|
||||
/// Create a new L0 cache with the specified capacity
|
||||
pub fn new(max_entries: usize) -> Self {
|
||||
Self {
|
||||
map: Arc::new(DashMap::with_capacity(max_entries.min(1024))),
|
||||
max_entries,
|
||||
hits: AtomicU64::new(0),
|
||||
misses: AtomicU64::new(0),
|
||||
}
|
||||
}
|
||||
|
||||
/// Get an encoding from the cache (returns Arc for zero-copy access)
|
||||
#[inline]
|
||||
pub fn get(&self, key: &str) -> Option<Arc<Encoding>> {
|
||||
match self.map.get(key) {
|
||||
Some(entry) => {
|
||||
self.hits.fetch_add(1, Ordering::Relaxed);
|
||||
// Arc::clone is cheap (just increment reference count)
|
||||
Some(Arc::clone(entry.value()))
|
||||
}
|
||||
None => {
|
||||
self.misses.fetch_add(1, Ordering::Relaxed);
|
||||
None
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Insert an encoding into the cache
|
||||
pub fn insert(&self, key: String, value: Encoding) {
|
||||
// Simple eviction: if we're at capacity, remove a random entry
|
||||
// DashMap doesn't support LRU directly, so we use a simple strategy
|
||||
if self.map.len() >= self.max_entries {
|
||||
let key_to_remove = { self.map.iter().next().map(|entry| entry.key().clone()) };
|
||||
|
||||
// Now remove it
|
||||
if let Some(k) = key_to_remove {
|
||||
self.map.remove(&k);
|
||||
}
|
||||
}
|
||||
|
||||
self.map.insert(key, Arc::new(value));
|
||||
}
|
||||
|
||||
/// Insert a pre-wrapped Arc encoding into the cache (avoids double-wrapping)
|
||||
pub fn insert_arc(&self, key: String, value: Arc<Encoding>) {
|
||||
if self.map.len() >= self.max_entries {
|
||||
let key_to_remove = { self.map.iter().next().map(|entry| entry.key().clone()) };
|
||||
if let Some(k) = key_to_remove {
|
||||
self.map.remove(&k);
|
||||
}
|
||||
}
|
||||
self.map.insert(key, value);
|
||||
}
|
||||
|
||||
/// Get the current number of entries in the cache
|
||||
pub fn len(&self) -> usize {
|
||||
self.map.len()
|
||||
}
|
||||
|
||||
/// Check if the cache is empty
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.map.is_empty()
|
||||
}
|
||||
|
||||
/// Get cache statistics
|
||||
pub fn stats(&self) -> CacheStats {
|
||||
let hits = self.hits.load(Ordering::Relaxed);
|
||||
let misses = self.misses.load(Ordering::Relaxed);
|
||||
let total_requests = hits + misses;
|
||||
|
||||
CacheStats {
|
||||
hits,
|
||||
misses,
|
||||
entries: self.len(),
|
||||
hit_rate: if total_requests > 0 {
|
||||
hits as f64 / total_requests as f64
|
||||
} else {
|
||||
0.0
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
/// Clear the cache
|
||||
pub fn clear(&self) {
|
||||
self.map.clear();
|
||||
self.hits.store(0, Ordering::Relaxed);
|
||||
self.misses.store(0, Ordering::Relaxed);
|
||||
}
|
||||
|
||||
/// Estimate memory usage in bytes
|
||||
pub fn memory_usage(&self) -> usize {
|
||||
// Rough estimate:
|
||||
// - Each entry: key (string) + value (encoding ~250 tokens * 4 bytes) + overhead
|
||||
// - Average: ~2.2KB per entry
|
||||
self.len() * 2200
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct CacheStats {
|
||||
pub hits: u64,
|
||||
pub misses: u64,
|
||||
pub entries: usize,
|
||||
pub hit_rate: f64,
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::tokenizer::traits::Encoding;
|
||||
|
||||
fn mock_encoding(tokens: Vec<u32>) -> Encoding {
|
||||
Encoding::Sp(tokens)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_basic_get_set() {
|
||||
let cache = L0Cache::new(10);
|
||||
|
||||
// Miss
|
||||
assert!(cache.get("hello").is_none());
|
||||
|
||||
// Insert
|
||||
cache.insert("hello".to_string(), mock_encoding(vec![1, 2, 3]));
|
||||
|
||||
// Hit - now returns Arc<Encoding>
|
||||
let result = cache.get("hello");
|
||||
assert!(result.is_some());
|
||||
assert_eq!(result.unwrap().token_ids(), &[1, 2, 3]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_eviction() {
|
||||
let cache = L0Cache::new(2);
|
||||
|
||||
cache.insert("a".to_string(), mock_encoding(vec![1]));
|
||||
cache.insert("b".to_string(), mock_encoding(vec![2]));
|
||||
|
||||
// Should evict when adding third
|
||||
cache.insert("c".to_string(), mock_encoding(vec![3]));
|
||||
|
||||
// Cache should have exactly 2 entries
|
||||
assert_eq!(cache.len(), 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_stats() {
|
||||
let cache = L0Cache::new(10);
|
||||
|
||||
cache.insert("test".to_string(), mock_encoding(vec![1, 2, 3]));
|
||||
|
||||
// 1 miss (initial get that returned None)
|
||||
let _ = cache.get("missing");
|
||||
|
||||
// 1 hit
|
||||
let _ = cache.get("test");
|
||||
|
||||
let stats = cache.stats();
|
||||
assert_eq!(stats.hits, 1);
|
||||
assert_eq!(stats.misses, 1);
|
||||
assert_eq!(stats.hit_rate, 0.5);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_clear() {
|
||||
let cache = L0Cache::new(10);
|
||||
|
||||
cache.insert("test".to_string(), mock_encoding(vec![1, 2, 3]));
|
||||
assert_eq!(cache.len(), 1);
|
||||
|
||||
cache.clear();
|
||||
assert_eq!(cache.len(), 0);
|
||||
assert!(cache.get("test").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_concurrent_access() {
|
||||
use std::thread;
|
||||
|
||||
let cache = Arc::new(L0Cache::new(1000));
|
||||
let mut handles = vec![];
|
||||
|
||||
// Spawn 10 threads
|
||||
for i in 0..10 {
|
||||
let cache_clone = cache.clone();
|
||||
handles.push(thread::spawn(move || {
|
||||
// Each thread inserts and reads
|
||||
let key = format!("key_{}", i);
|
||||
cache_clone.insert(key.clone(), mock_encoding(vec![i as u32]));
|
||||
|
||||
// Read it back
|
||||
let result = cache_clone.get(&key);
|
||||
assert!(result.is_some());
|
||||
}));
|
||||
}
|
||||
|
||||
for handle in handles {
|
||||
handle.join().unwrap();
|
||||
}
|
||||
|
||||
// Should have 10 entries
|
||||
assert_eq!(cache.len(), 10);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_arc_reuse() {
|
||||
// Test that multiple gets return the same Arc (reference counting)
|
||||
let cache = L0Cache::new(10);
|
||||
cache.insert("test".to_string(), mock_encoding(vec![1, 2, 3]));
|
||||
|
||||
let arc1 = cache.get("test").unwrap();
|
||||
let arc2 = cache.get("test").unwrap();
|
||||
|
||||
// Both should point to the same allocation
|
||||
assert!(Arc::ptr_eq(&arc1, &arc2));
|
||||
}
|
||||
}
|
||||
524
sgl-model-gateway/src/tokenizer/cache/l1.rs
vendored
524
sgl-model-gateway/src/tokenizer/cache/l1.rs
vendored
@@ -1,524 +0,0 @@
|
||||
//! L1 Cache: Special-token boundary prefix cache
|
||||
//!
|
||||
//! Caches tokenization results at ALL special token boundaries.
|
||||
//! Special tokens (like `<|im_start|>`, `<|im_end|>`) are atomic in BPE tokenizers (special: true, normalized: false),
|
||||
//! making them the ONLY safe split points that guarantee correctness.
|
||||
//!
|
||||
//! **Design**: Cache at every special token boundary (not at fixed granularity intervals)
|
||||
//! - Simple: No granularity parameter, no search windows
|
||||
//! - Efficient: Fewer cache entries (10 instead of 64 for typical 8KB prompt)
|
||||
//! - Natural: Aligns with actual chat template structure
|
||||
//!
|
||||
//! Example:
|
||||
//!
|
||||
//! Template: "<|im_start|>system\nYou are helpful.<|im_end|><|im_start|>user\n{query}<|im_end|>"
|
||||
//!
|
||||
//! Request 1: "<|im_start|>system\nYou are helpful.<|im_end|><|im_start|>user\nWhat is 2+2?<|im_end|>"
|
||||
//! Request 2: "<|im_start|>system\nYou are helpful.<|im_end|><|im_start|>user\nHello!<|im_end|>"
|
||||
//!
|
||||
//! Cache points: After each "<|im_end|>" (atomic tokens, guaranteed safe)
|
||||
//! Result: tokenize(prefix) + tokenize(suffix) == tokenize(prefix + suffix)
|
||||
|
||||
use std::{
|
||||
mem::size_of,
|
||||
sync::{
|
||||
atomic::{AtomicU64, Ordering},
|
||||
Arc,
|
||||
},
|
||||
};
|
||||
|
||||
use blake3;
|
||||
use dashmap::DashMap;
|
||||
|
||||
use super::super::traits::TokenIdType;
|
||||
|
||||
/// Hash type for cache keys
|
||||
type Blake3Hash = [u8; 32];
|
||||
|
||||
/// Number of shards for concurrent access
|
||||
const NUM_SHARDS: usize = 16;
|
||||
|
||||
/// Find ALL special token boundaries in the text
|
||||
///
|
||||
/// **ONLY uses special tokens** - these are atomic (special: true, normalized: false) in BPE,
|
||||
/// guaranteeing: tokenize(prefix) + tokenize(suffix) == tokenize(prefix + suffix)
|
||||
///
|
||||
/// No fallback to whitespace/punctuation - better to not cache than risk corruption.
|
||||
///
|
||||
/// Common special tokens:
|
||||
/// - ChatML: `<|im_start|>`, `<|im_end|>`
|
||||
/// - Llama 3: `<|begin_of_text|>`, `<|end_of_text|>`, `<|eot_id|>`
|
||||
/// - GPT: `<|endoftext|>`
|
||||
/// - Custom: `<|reserved_special_token_N|>`
|
||||
///
|
||||
/// Returns positions immediately after each special token (where prefixes can be cached).
|
||||
fn find_special_token_boundaries(text: &str, special_tokens: &[&str]) -> Vec<usize> {
|
||||
if special_tokens.is_empty() {
|
||||
return Vec::new();
|
||||
}
|
||||
|
||||
let mut boundaries = Vec::new();
|
||||
|
||||
// Find all special token end positions
|
||||
for &token in special_tokens {
|
||||
let mut start = 0;
|
||||
while let Some(pos) = text[start..].find(token) {
|
||||
let boundary = start + pos + token.len();
|
||||
// Only cache boundaries that leave some suffix to tokenize
|
||||
if boundary < text.len() {
|
||||
boundaries.push(boundary);
|
||||
}
|
||||
start = boundary;
|
||||
}
|
||||
}
|
||||
|
||||
// Sort and deduplicate (in case multiple special tokens end at same position)
|
||||
boundaries.sort_unstable();
|
||||
boundaries.dedup();
|
||||
|
||||
boundaries
|
||||
}
|
||||
|
||||
/// A cached prefix entry
|
||||
/// Uses Arc<[TokenIdType]> for zero-copy access to tokens
|
||||
#[derive(Debug, Clone)]
|
||||
struct CachedPrefix {
|
||||
/// The pre-computed token IDs for this prefix (Arc for zero-copy cloning)
|
||||
tokens: Arc<[TokenIdType]>,
|
||||
/// Last access timestamp (for LRU eviction)
|
||||
last_accessed: Arc<AtomicU64>,
|
||||
/// Size in bytes (for memory tracking during eviction)
|
||||
size_bytes: usize,
|
||||
}
|
||||
|
||||
/// L1 cache implementation with special-token-boundary prefix matching
|
||||
pub struct L1Cache {
|
||||
/// Sharded maps for concurrent access
|
||||
/// Key: Blake3 hash of bytes[0..boundary]
|
||||
/// Value: Cached token IDs for that prefix
|
||||
shards: Vec<Arc<DashMap<Blake3Hash, CachedPrefix>>>,
|
||||
/// Maximum memory in bytes
|
||||
max_memory: usize,
|
||||
/// Current memory usage estimate
|
||||
current_memory: AtomicU64,
|
||||
/// Cache hit counter
|
||||
hits: AtomicU64,
|
||||
/// Cache miss counter
|
||||
misses: AtomicU64,
|
||||
/// Monotonic counter for LRU timestamps
|
||||
access_counter: AtomicU64,
|
||||
}
|
||||
|
||||
impl L1Cache {
|
||||
/// Create a new L1 cache with the specified memory limit
|
||||
pub fn new(max_memory: usize) -> Self {
|
||||
let shards = (0..NUM_SHARDS).map(|_| Arc::new(DashMap::new())).collect();
|
||||
|
||||
Self {
|
||||
shards,
|
||||
max_memory,
|
||||
current_memory: AtomicU64::new(0),
|
||||
hits: AtomicU64::new(0),
|
||||
misses: AtomicU64::new(0),
|
||||
access_counter: AtomicU64::new(0),
|
||||
}
|
||||
}
|
||||
|
||||
/// Try to find the longest prefix match at special token boundaries
|
||||
/// Returns (cached_tokens, byte_offset) if found
|
||||
///
|
||||
/// Uses pre-computed tokens cached during insertion.
|
||||
/// Returns Vec<TokenIdType> as the caller needs to extend it with suffix tokens.
|
||||
pub fn longest_prefix_match(
|
||||
&self,
|
||||
input: &str,
|
||||
special_tokens: &[&str],
|
||||
) -> Option<(Vec<TokenIdType>, usize)> {
|
||||
let boundaries = find_special_token_boundaries(input, special_tokens);
|
||||
|
||||
if boundaries.is_empty() {
|
||||
self.misses.fetch_add(1, Ordering::Relaxed);
|
||||
return None;
|
||||
}
|
||||
|
||||
// Build all prefix hashes incrementally O(N)
|
||||
let mut hasher = blake3::Hasher::new();
|
||||
let mut prefix_hashes = Vec::with_capacity(boundaries.len());
|
||||
let mut last_pos = 0;
|
||||
let bytes = input.as_bytes();
|
||||
for &boundary_pos in &boundaries {
|
||||
hasher.update(&bytes[last_pos..boundary_pos]);
|
||||
prefix_hashes.push((boundary_pos, *hasher.clone().finalize().as_bytes()));
|
||||
last_pos = boundary_pos;
|
||||
}
|
||||
|
||||
// Search from the longest boundary to find the best match
|
||||
for (boundary_pos, hash_bytes) in prefix_hashes.into_iter().rev() {
|
||||
let shard_idx = hash_bytes[0] as usize % NUM_SHARDS;
|
||||
|
||||
if let Some(entry) = self.shards[shard_idx].get(&hash_bytes) {
|
||||
// Update last accessed timestamp for LRU
|
||||
let timestamp = self.access_counter.fetch_add(1, Ordering::Relaxed);
|
||||
entry.last_accessed.store(timestamp, Ordering::Relaxed);
|
||||
|
||||
self.hits.fetch_add(1, Ordering::Relaxed);
|
||||
// Convert Arc<[T]> to Vec<T> - caller will extend with suffix tokens
|
||||
return Some((entry.tokens.to_vec(), boundary_pos));
|
||||
}
|
||||
}
|
||||
|
||||
self.misses.fetch_add(1, Ordering::Relaxed);
|
||||
None
|
||||
}
|
||||
|
||||
/// Insert prefix entries at ALL special token boundaries
|
||||
///
|
||||
/// Uses incremental hashing and tokenization for O(N) performance.
|
||||
///
|
||||
/// Optimized for workloads with high prefix reuse (e.g., chat templates with repeated system prompts).
|
||||
pub fn insert_at_boundaries<E: super::super::traits::Encoder + ?Sized>(
|
||||
&self,
|
||||
input: &str,
|
||||
tokenizer: &E,
|
||||
special_tokens: &[&str],
|
||||
add_special_tokens: bool,
|
||||
) -> anyhow::Result<()> {
|
||||
let boundaries = find_special_token_boundaries(input, special_tokens);
|
||||
|
||||
if boundaries.is_empty() {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let mut hasher = blake3::Hasher::new();
|
||||
let mut running_tokens = Vec::new();
|
||||
let mut last_pos = 0;
|
||||
let mut entries_to_insert = Vec::with_capacity(boundaries.len());
|
||||
let bytes = input.as_bytes();
|
||||
for (i, &boundary_pos) in boundaries.iter().enumerate() {
|
||||
let delta_text = &input[last_pos..boundary_pos];
|
||||
|
||||
// 1. Incremental Hash update
|
||||
hasher.update(&bytes[last_pos..boundary_pos]);
|
||||
let hash_bytes: Blake3Hash = *hasher.clone().finalize().as_bytes();
|
||||
|
||||
// 2. Incremental Tokenization
|
||||
// Only add special tokens (like BOS) for the very first segment to avoid duplicates
|
||||
let segment_encoding = tokenizer.encode(delta_text, (i == 0) && add_special_tokens)?;
|
||||
running_tokens.extend_from_slice(segment_encoding.token_ids());
|
||||
|
||||
// 3. Prepare entry
|
||||
// Convert current tokens to Arc<[TokenIdType]> for sharing
|
||||
let prefix_tokens: Arc<[TokenIdType]> = running_tokens.as_slice().into();
|
||||
|
||||
// Size = text bytes + token storage
|
||||
let size_bytes = boundary_pos + prefix_tokens.len() * size_of::<TokenIdType>();
|
||||
|
||||
entries_to_insert.push((hash_bytes, prefix_tokens, size_bytes));
|
||||
|
||||
last_pos = boundary_pos;
|
||||
}
|
||||
|
||||
if entries_to_insert.is_empty() {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let total_size_needed: usize = entries_to_insert.iter().map(|(_, _, size)| size).sum();
|
||||
|
||||
// Evict if necessary
|
||||
let current = self.current_memory.load(Ordering::Relaxed) as usize;
|
||||
if current + total_size_needed > self.max_memory {
|
||||
self.evict_lru(total_size_needed);
|
||||
}
|
||||
|
||||
// Insert all entries
|
||||
let current_timestamp = self.access_counter.load(Ordering::Relaxed);
|
||||
for (hash_bytes, prefix_tokens, size_bytes) in entries_to_insert {
|
||||
let shard_idx = hash_bytes[0] as usize % NUM_SHARDS;
|
||||
|
||||
let cached = CachedPrefix {
|
||||
tokens: prefix_tokens, // Already Arc<[TokenIdType]>
|
||||
last_accessed: Arc::new(AtomicU64::new(current_timestamp)),
|
||||
size_bytes,
|
||||
};
|
||||
|
||||
self.shards[shard_idx].insert(hash_bytes, cached);
|
||||
self.current_memory
|
||||
.fetch_add(size_bytes as u64, Ordering::Relaxed);
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Evict least recently used entries using approximate LRU via random sampling
|
||||
///
|
||||
/// This uses an approximate LRU strategy that's much faster than true LRU:
|
||||
/// - Samples K random entries from the cache (K=32)
|
||||
/// - Evicts the oldest entry among the samples
|
||||
/// - Repeats until enough space is freed
|
||||
///
|
||||
/// This provides O(samples) complexity instead of O(total_entries * log(total_entries)),
|
||||
/// avoiding latency spikes when eviction is triggered on large caches.
|
||||
///
|
||||
/// The approximation is excellent in practice - sampling 32 entries from a large cache
|
||||
/// gives high probability of finding very old entries.
|
||||
fn evict_lru(&self, space_needed: usize) {
|
||||
const SAMPLE_SIZE: usize = 32; // Number of entries to sample per eviction round
|
||||
let mut freed = 0usize;
|
||||
let mut iteration = 0usize;
|
||||
|
||||
// Keep evicting until we have enough space
|
||||
while freed < space_needed {
|
||||
// Collect samples from shards
|
||||
let mut samples: Vec<(usize, Blake3Hash, u64, usize)> = Vec::with_capacity(SAMPLE_SIZE);
|
||||
|
||||
// Sample entries across different shards
|
||||
for i in 0..SAMPLE_SIZE {
|
||||
// Distribute samples across shards using iteration and index for variety
|
||||
let shard_idx = (iteration * SAMPLE_SIZE + i) % NUM_SHARDS;
|
||||
|
||||
// Get first entry from that shard (DashMap iteration order is arbitrary)
|
||||
if let Some(entry) = self.shards[shard_idx].iter().next() {
|
||||
let hash = *entry.key();
|
||||
let timestamp = entry.value().last_accessed.load(Ordering::Relaxed);
|
||||
let size = entry.value().size_bytes;
|
||||
samples.push((shard_idx, hash, timestamp, size));
|
||||
}
|
||||
}
|
||||
|
||||
if samples.is_empty() {
|
||||
// Cache is empty, nothing to evict
|
||||
break;
|
||||
}
|
||||
|
||||
// Find the oldest entry among samples
|
||||
if let Some((shard_idx, hash, _, _)) =
|
||||
samples.iter().min_by_key(|(_, _, ts, _)| ts).copied()
|
||||
{
|
||||
// Remove it
|
||||
if let Some((_, removed)) = self.shards[shard_idx].remove(&hash) {
|
||||
freed += removed.size_bytes;
|
||||
self.current_memory
|
||||
.fetch_sub(removed.size_bytes as u64, Ordering::Relaxed);
|
||||
}
|
||||
}
|
||||
|
||||
iteration += 1;
|
||||
}
|
||||
}
|
||||
|
||||
/// Get the number of entries in the cache
|
||||
pub fn len(&self) -> usize {
|
||||
self.shards.iter().map(|s| s.len()).sum()
|
||||
}
|
||||
|
||||
/// Check if the cache is empty
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.shards.iter().all(|s| s.is_empty())
|
||||
}
|
||||
|
||||
/// Get cache statistics
|
||||
pub fn stats(&self) -> L1CacheStats {
|
||||
let hits = self.hits.load(Ordering::Relaxed);
|
||||
let misses = self.misses.load(Ordering::Relaxed);
|
||||
let total_requests = hits + misses;
|
||||
|
||||
L1CacheStats {
|
||||
hits,
|
||||
misses,
|
||||
entries: self.len(),
|
||||
memory_bytes: self.current_memory.load(Ordering::Relaxed) as usize,
|
||||
hit_rate: if total_requests > 0 {
|
||||
hits as f64 / total_requests as f64
|
||||
} else {
|
||||
0.0
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
/// Clear the cache
|
||||
pub fn clear(&self) {
|
||||
for shard in &self.shards {
|
||||
shard.clear();
|
||||
}
|
||||
self.current_memory.store(0, Ordering::Relaxed);
|
||||
self.hits.store(0, Ordering::Relaxed);
|
||||
self.misses.store(0, Ordering::Relaxed);
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct L1CacheStats {
|
||||
pub hits: u64,
|
||||
pub misses: u64,
|
||||
pub entries: usize,
|
||||
pub memory_bytes: usize,
|
||||
pub hit_rate: f64,
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::tokenizer::mock::MockTokenizer;
|
||||
|
||||
#[test]
|
||||
fn test_basic_prefix_match() {
|
||||
let cache = L1Cache::new(1024 * 1024);
|
||||
let special_tokens = &["<|im_start|>", "<|im_end|>"];
|
||||
let tokenizer = MockTokenizer::new();
|
||||
|
||||
// Realistic ChatML template with special tokens
|
||||
let input1 = "<|im_start|>system\nYou are a helpful assistant that provides clear and detailed responses.<|im_end|><|im_start|>user\nHello there! How are you doing today?<|im_end|>";
|
||||
|
||||
// Insert at special token boundaries (re-tokenizes prefixes)
|
||||
cache
|
||||
.insert_at_boundaries(input1, &tokenizer, special_tokens, false)
|
||||
.unwrap();
|
||||
|
||||
// Should have cached at special token boundaries
|
||||
assert!(!cache.is_empty());
|
||||
|
||||
// Search with same prefix but different user query
|
||||
let input2 = "<|im_start|>system\nYou are a helpful assistant that provides clear and detailed responses.<|im_end|><|im_start|>user\nWhat is 2+2?<|im_end|>";
|
||||
let result = cache.longest_prefix_match(input2, special_tokens);
|
||||
|
||||
// Should find a match at the special token boundary (after system message)
|
||||
assert!(result.is_some());
|
||||
let (tokens, offset) = result.unwrap();
|
||||
assert!(offset > 0);
|
||||
assert!(!tokens.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_short_input_with_boundaries() {
|
||||
let cache = L1Cache::new(1024 * 1024);
|
||||
let special_tokens = &["<|im_start|>", "<|im_end|>"];
|
||||
let tokenizer = MockTokenizer::new();
|
||||
|
||||
// Short input with special tokens
|
||||
let input = "<|im_start|>user\nHi<|im_end|>";
|
||||
|
||||
cache
|
||||
.insert_at_boundaries(input, &tokenizer, special_tokens, false)
|
||||
.unwrap();
|
||||
|
||||
// Should cache at <|im_start|> boundary (has suffix left)
|
||||
assert!(!cache.is_empty());
|
||||
|
||||
// Should find a match
|
||||
let result = cache.longest_prefix_match(input, special_tokens);
|
||||
assert!(result.is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_longest_match() {
|
||||
let cache = L1Cache::new(1024 * 1024);
|
||||
let special_tokens = &["<|im_start|>", "<|im_end|>"];
|
||||
let tokenizer = MockTokenizer::new();
|
||||
|
||||
// Create multi-turn conversation with multiple special token boundaries (~400 bytes)
|
||||
let input = "<|im_start|>system\nYou are a helpful AI assistant that provides detailed and accurate responses.<|im_end|><|im_start|>user\nHello there! How are you today? Can you help me understand how tokenization works in language models?<|im_end|><|im_start|>assistant\nI'm doing well, thank you! I'd be happy to explain tokenization. Tokenization is the process of breaking text into smaller units called tokens.<|im_end|>";
|
||||
|
||||
cache
|
||||
.insert_at_boundaries(input, &tokenizer, special_tokens, false)
|
||||
.unwrap();
|
||||
|
||||
// Should have multiple entries at special token boundaries
|
||||
assert!(cache.len() >= 2); // At least 2 boundaries
|
||||
|
||||
// Search with partial conversation - should match at a special token boundary
|
||||
let partial_input = "<|im_start|>system\nYou are a helpful AI assistant that provides detailed and accurate responses.<|im_end|><|im_start|>user\nHello there! How are you today? Can you help me understand how tokenization works in language models?<|im_end|>";
|
||||
let result = cache.longest_prefix_match(partial_input, special_tokens);
|
||||
|
||||
// Should find a match at a special token boundary
|
||||
assert!(result.is_some());
|
||||
let (_, offset) = result.unwrap();
|
||||
assert!(offset > 0);
|
||||
assert!(offset <= partial_input.len());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_stats() {
|
||||
let cache = L1Cache::new(1024 * 1024);
|
||||
let special_tokens = &["<|im_start|>", "<|im_end|>"];
|
||||
let tokenizer = MockTokenizer::new();
|
||||
|
||||
// ChatML input with special tokens
|
||||
let input = "<|im_start|>system\nYou are a helpful assistant that provides detailed answers.<|im_end|><|im_start|>user\nHello there! How are you today?<|im_end|>";
|
||||
|
||||
cache
|
||||
.insert_at_boundaries(input, &tokenizer, special_tokens, false)
|
||||
.unwrap();
|
||||
|
||||
// Try to find match
|
||||
let _ = cache.longest_prefix_match(input, special_tokens);
|
||||
|
||||
let stats = cache.stats();
|
||||
// Should have at least one hit (the longest special token boundary should match)
|
||||
assert!(stats.hits >= 1);
|
||||
assert_eq!(stats.hit_rate, 1.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_clear() {
|
||||
let cache = L1Cache::new(1024 * 1024);
|
||||
let special_tokens = &["<|im_start|>", "<|im_end|>"];
|
||||
let tokenizer = MockTokenizer::new();
|
||||
|
||||
// ChatML input with special tokens
|
||||
let input = "<|im_start|>system\nYou are a helpful assistant that provides clear and detailed responses.<|im_end|><|im_start|>user\nHello there!<|im_end|>";
|
||||
|
||||
cache
|
||||
.insert_at_boundaries(input, &tokenizer, special_tokens, false)
|
||||
.unwrap();
|
||||
assert!(!cache.is_empty());
|
||||
|
||||
cache.clear();
|
||||
assert!(cache.is_empty());
|
||||
|
||||
let stats = cache.stats();
|
||||
assert_eq!(stats.hits, 0);
|
||||
assert_eq!(stats.misses, 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_lru_eviction() {
|
||||
// Create a small cache (5KB) to trigger eviction
|
||||
let cache = L1Cache::new(5 * 1024);
|
||||
let special_tokens = &["<|im_start|>", "<|im_end|>", "<|eot_id|>"];
|
||||
let tokenizer = MockTokenizer::new();
|
||||
|
||||
// Insert first conversation
|
||||
let input1 = "<|im_start|>system\nYou are a helpful assistant specialized in mathematics.<|im_end|><|im_start|>user\nCan you explain calculus to me?<|im_end|><|im_start|>assistant\nCertainly! Calculus is a branch of mathematics that studies continuous change.<|im_end|><|eot_id|>";
|
||||
cache
|
||||
.insert_at_boundaries(input1, &tokenizer, special_tokens, false)
|
||||
.unwrap();
|
||||
|
||||
// Access the first entry to update its timestamp
|
||||
let result = cache.longest_prefix_match(input1, special_tokens);
|
||||
assert!(result.is_some());
|
||||
|
||||
// Insert second conversation
|
||||
let input2 = "<|im_start|>system\nYou are a helpful assistant specialized in physics.<|im_end|><|im_start|>user\nWhat is quantum mechanics?<|im_end|><|im_start|>assistant\nQuantum mechanics is the fundamental theory describing nature at atomic and subatomic scales.<|im_end|><|eot_id|>";
|
||||
cache
|
||||
.insert_at_boundaries(input2, &tokenizer, special_tokens, false)
|
||||
.unwrap();
|
||||
|
||||
// Access the second entry to make it more recent
|
||||
let result = cache.longest_prefix_match(input2, special_tokens);
|
||||
assert!(result.is_some());
|
||||
|
||||
// Insert third conversation (should trigger eviction of oldest)
|
||||
let input3 = "<|im_start|>system\nYou are a helpful assistant specialized in chemistry.<|im_end|><|im_start|>user\nExplain the periodic table to me please.<|im_end|><|im_start|>assistant\nThe periodic table is a tabular arrangement of chemical elements organized by atomic number and electron configuration.<|im_end|><|eot_id|>";
|
||||
cache
|
||||
.insert_at_boundaries(input3, &tokenizer, special_tokens, false)
|
||||
.unwrap();
|
||||
|
||||
// Verify cache didn't exceed max memory
|
||||
let stats = cache.stats();
|
||||
assert!(stats.memory_bytes <= 5 * 1024);
|
||||
|
||||
// The most recently accessed entries should still be present
|
||||
let result = cache.longest_prefix_match(input3, special_tokens);
|
||||
assert!(result.is_some());
|
||||
}
|
||||
}
|
||||
372
sgl-model-gateway/src/tokenizer/cache/mod.rs
vendored
372
sgl-model-gateway/src/tokenizer/cache/mod.rs
vendored
@@ -1,372 +0,0 @@
|
||||
//! Tokenizer Caching Layer
|
||||
//!
|
||||
//! Provides a caching wrapper around any tokenizer implementation to speed up
|
||||
//! repeated tokenization of the same strings (e.g., system prompts).
|
||||
//!
|
||||
//! # Architecture
|
||||
//! - **L0 Cache**: Whole-string exact match (90% of wins)
|
||||
//! - **L1 Cache**: Prefix matching at fixed boundaries (future work)
|
||||
//!
|
||||
//! # Usage
|
||||
//! ```ignore
|
||||
//! let tokenizer = Arc::new(HuggingFaceTokenizer::from_file("tokenizer.json")?);
|
||||
//! let cached = Arc::new(CachedTokenizer::new(tokenizer, CacheConfig::default()));
|
||||
//! let encoding = cached.encode("Hello world")?;
|
||||
//! ```
|
||||
|
||||
mod fingerprint;
|
||||
mod l0;
|
||||
mod l1;
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use anyhow::Result;
|
||||
pub use fingerprint::TokenizerFingerprint;
|
||||
pub use l0::{CacheStats, L0Cache};
|
||||
pub use l1::{L1Cache, L1CacheStats};
|
||||
use rayon::prelude::*;
|
||||
|
||||
use super::traits::{Decoder, Encoder, Encoding, SpecialTokens, TokenIdType, Tokenizer};
|
||||
|
||||
/// Configuration for the tokenizer cache
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct CacheConfig {
|
||||
/// Enable L0 (whole-string) cache
|
||||
pub enable_l0: bool,
|
||||
/// Maximum number of entries in L0 cache
|
||||
pub l0_max_entries: usize,
|
||||
/// Enable L1 (prefix) cache
|
||||
pub enable_l1: bool,
|
||||
/// Maximum memory for L1 cache in bytes
|
||||
pub l1_max_memory: usize,
|
||||
}
|
||||
|
||||
impl Default for CacheConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
enable_l0: true,
|
||||
l0_max_entries: 10_000, // ~22MB memory for typical prompts
|
||||
enable_l1: false, // Opt-in for now
|
||||
l1_max_memory: 50 * 1024 * 1024, // 50MB
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// A caching wrapper around any tokenizer
|
||||
pub struct CachedTokenizer {
|
||||
/// The underlying tokenizer
|
||||
inner: Arc<dyn Tokenizer>,
|
||||
/// L0 cache (whole-string exact match)
|
||||
l0: Option<L0Cache>,
|
||||
/// L1 cache (prefix matching at fixed boundaries)
|
||||
l1: Option<L1Cache>,
|
||||
/// Configuration
|
||||
#[allow(dead_code)]
|
||||
config: CacheConfig,
|
||||
/// Fingerprint for cache invalidation
|
||||
fingerprint: TokenizerFingerprint,
|
||||
/// Cached special token strings (extracted once at construction)
|
||||
special_token_strings: Vec<String>,
|
||||
}
|
||||
|
||||
impl CachedTokenizer {
|
||||
/// Create a new cached tokenizer
|
||||
pub fn new(inner: Arc<dyn Tokenizer>, config: CacheConfig) -> Self {
|
||||
let fingerprint = TokenizerFingerprint::from_tokenizer(inner.as_ref());
|
||||
|
||||
let l0 = if config.enable_l0 {
|
||||
Some(L0Cache::new(config.l0_max_entries))
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
let l1 = if config.enable_l1 {
|
||||
Some(L1Cache::new(config.l1_max_memory))
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
// Extract special tokens once at construction time
|
||||
let special_token_strings = Self::extract_special_token_strings(&inner);
|
||||
|
||||
Self {
|
||||
inner,
|
||||
l0,
|
||||
l1,
|
||||
config,
|
||||
fingerprint,
|
||||
special_token_strings,
|
||||
}
|
||||
}
|
||||
|
||||
/// Extract all special token strings from the tokenizer (called once at construction)
|
||||
fn extract_special_token_strings(tokenizer: &Arc<dyn Tokenizer>) -> Vec<String> {
|
||||
let special_tokens = tokenizer.get_special_tokens();
|
||||
let mut tokens = Vec::new();
|
||||
|
||||
if let Some(ref token) = special_tokens.bos_token {
|
||||
tokens.push(token.clone());
|
||||
}
|
||||
if let Some(ref token) = special_tokens.eos_token {
|
||||
tokens.push(token.clone());
|
||||
}
|
||||
if let Some(ref token) = special_tokens.unk_token {
|
||||
tokens.push(token.clone());
|
||||
}
|
||||
if let Some(ref token) = special_tokens.sep_token {
|
||||
tokens.push(token.clone());
|
||||
}
|
||||
if let Some(ref token) = special_tokens.pad_token {
|
||||
tokens.push(token.clone());
|
||||
}
|
||||
if let Some(ref token) = special_tokens.cls_token {
|
||||
tokens.push(token.clone());
|
||||
}
|
||||
if let Some(ref token) = special_tokens.mask_token {
|
||||
tokens.push(token.clone());
|
||||
}
|
||||
|
||||
tokens.extend(special_tokens.additional_special_tokens.iter().cloned());
|
||||
tokens
|
||||
}
|
||||
|
||||
/// Get L0 cache statistics
|
||||
pub fn cache_stats(&self) -> Option<CacheStats> {
|
||||
self.l0.as_ref().map(|cache| cache.stats())
|
||||
}
|
||||
|
||||
/// Get L1 cache statistics
|
||||
pub fn l1_cache_stats(&self) -> Option<L1CacheStats> {
|
||||
self.l1.as_ref().map(|cache| cache.stats())
|
||||
}
|
||||
|
||||
/// Clear the cache
|
||||
pub fn clear_cache(&self) {
|
||||
if let Some(l0) = &self.l0 {
|
||||
l0.clear();
|
||||
}
|
||||
if let Some(l1) = &self.l1 {
|
||||
l1.clear();
|
||||
}
|
||||
}
|
||||
|
||||
/// Get the fingerprint of the underlying tokenizer
|
||||
pub fn fingerprint(&self) -> &TokenizerFingerprint {
|
||||
&self.fingerprint
|
||||
}
|
||||
|
||||
/// Get a reference to the inner (wrapped) tokenizer
|
||||
pub fn inner(&self) -> &Arc<dyn Tokenizer> {
|
||||
&self.inner
|
||||
}
|
||||
}
|
||||
|
||||
impl Encoder for CachedTokenizer {
|
||||
fn encode(&self, input: &str, add_special_tokens: bool) -> Result<Encoding> {
|
||||
// L0 cache lookup (exact match) - returns Arc<Encoding> for zero-copy
|
||||
// Note: L0 cache doesn't distinguish by add_special_tokens flag
|
||||
// This is acceptable for the current use case where embeddings always use true
|
||||
// and chat always uses false with different input content
|
||||
if let Some(l0) = &self.l0 {
|
||||
if let Some(cached) = l0.get(input) {
|
||||
// Unwrap the Arc - since Encoding is Clone, we can return the inner value
|
||||
// For callers who need the tokens, they can access via token_ids() which is &[u32]
|
||||
return Ok((*cached).clone());
|
||||
}
|
||||
}
|
||||
|
||||
// L1 cache lookup (prefix match at special token boundaries)
|
||||
if let Some(l1) = &self.l1 {
|
||||
// Use pre-computed special tokens refs (avoids allocation per call)
|
||||
let tokens: Vec<&str> = self
|
||||
.special_token_strings
|
||||
.iter()
|
||||
.map(|s| s.as_str())
|
||||
.collect();
|
||||
|
||||
if let Some((prefix_tokens, prefix_len)) = l1.longest_prefix_match(input, &tokens) {
|
||||
// We have a prefix match - tokenize the suffix
|
||||
let suffix = &input[prefix_len..];
|
||||
if !suffix.is_empty() {
|
||||
let suffix_encoding = self.inner.encode(suffix, add_special_tokens)?;
|
||||
|
||||
// Merge prefix tokens + suffix tokens
|
||||
// Safe because we're splitting at special token boundaries
|
||||
let mut merged_tokens = prefix_tokens;
|
||||
merged_tokens.extend_from_slice(suffix_encoding.token_ids());
|
||||
|
||||
let merged_encoding = Encoding::Sp(merged_tokens);
|
||||
|
||||
// Cache the full result in L0
|
||||
if let Some(l0) = &self.l0 {
|
||||
l0.insert(input.to_string(), merged_encoding.clone());
|
||||
}
|
||||
|
||||
return Ok(merged_encoding);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Full tokenization (both L0 and L1 miss)
|
||||
let encoding = self.inner.encode(input, add_special_tokens)?;
|
||||
|
||||
// Cache in L0
|
||||
if let Some(l0) = &self.l0 {
|
||||
l0.insert(input.to_string(), encoding.clone());
|
||||
}
|
||||
|
||||
// Cache in L1 at special token boundaries
|
||||
// Re-tokenizes prefixes for correctness (optimized for high prefix reuse)
|
||||
if let Some(l1) = &self.l1 {
|
||||
let tokens: Vec<&str> = self
|
||||
.special_token_strings
|
||||
.iter()
|
||||
.map(|s| s.as_str())
|
||||
.collect();
|
||||
let _ =
|
||||
l1.insert_at_boundaries(input, self.inner.as_ref(), &tokens, add_special_tokens);
|
||||
// Ignore errors in cache insertion - cache is best-effort
|
||||
}
|
||||
|
||||
Ok(encoding)
|
||||
}
|
||||
|
||||
fn encode_batch(&self, inputs: &[&str], add_special_tokens: bool) -> Result<Vec<Encoding>> {
|
||||
// Process each input in parallel, leveraging thread-safe caches
|
||||
// This maintains the parallelism from the underlying HuggingFaceTokenizer
|
||||
inputs
|
||||
.par_iter()
|
||||
.map(|&input| self.encode(input, add_special_tokens))
|
||||
.collect()
|
||||
}
|
||||
}
|
||||
|
||||
impl Decoder for CachedTokenizer {
|
||||
fn decode(&self, token_ids: &[TokenIdType], skip_special_tokens: bool) -> Result<String> {
|
||||
// Decoding is not cached (it's fast enough and rarely repeated)
|
||||
self.inner.decode(token_ids, skip_special_tokens)
|
||||
}
|
||||
}
|
||||
|
||||
impl Tokenizer for CachedTokenizer {
|
||||
fn vocab_size(&self) -> usize {
|
||||
self.inner.vocab_size()
|
||||
}
|
||||
|
||||
fn get_special_tokens(&self) -> &SpecialTokens {
|
||||
self.inner.get_special_tokens()
|
||||
}
|
||||
|
||||
fn token_to_id(&self, token: &str) -> Option<TokenIdType> {
|
||||
self.inner.token_to_id(token)
|
||||
}
|
||||
|
||||
fn id_to_token(&self, id: TokenIdType) -> Option<String> {
|
||||
self.inner.id_to_token(id)
|
||||
}
|
||||
|
||||
fn as_any(&self) -> &dyn std::any::Any {
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::tokenizer::mock::MockTokenizer;
|
||||
|
||||
#[test]
|
||||
fn test_cache_hit() {
|
||||
let tokenizer = Arc::new(MockTokenizer::new());
|
||||
let cached = CachedTokenizer::new(tokenizer, CacheConfig::default());
|
||||
|
||||
let input = "Hello world";
|
||||
|
||||
// First call - miss
|
||||
let result1 = cached.encode(input, false).unwrap();
|
||||
|
||||
// Second call - hit
|
||||
let result2 = cached.encode(input, false).unwrap();
|
||||
|
||||
// Results should be identical
|
||||
assert_eq!(result1.token_ids(), result2.token_ids());
|
||||
|
||||
// Check cache stats
|
||||
let stats = cached.cache_stats().unwrap();
|
||||
assert_eq!(stats.hits, 1);
|
||||
assert_eq!(stats.misses, 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_cache_disabled() {
|
||||
let tokenizer = Arc::new(MockTokenizer::new());
|
||||
let config = CacheConfig {
|
||||
enable_l0: false,
|
||||
l0_max_entries: 0,
|
||||
enable_l1: false,
|
||||
l1_max_memory: 0,
|
||||
};
|
||||
let cached = CachedTokenizer::new(tokenizer, config);
|
||||
|
||||
let input = "Hello world";
|
||||
|
||||
// Both calls should work even without cache
|
||||
let result1 = cached.encode(input, false).unwrap();
|
||||
let result2 = cached.encode(input, false).unwrap();
|
||||
|
||||
assert_eq!(result1.token_ids(), result2.token_ids());
|
||||
|
||||
// No cache stats available
|
||||
assert!(cached.cache_stats().is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_encode_batch() {
|
||||
let tokenizer = Arc::new(MockTokenizer::new());
|
||||
let cached = CachedTokenizer::new(tokenizer, CacheConfig::default());
|
||||
|
||||
let inputs = vec!["Hello", "world", "Hello"]; // "Hello" repeated
|
||||
|
||||
let results = cached.encode_batch(&inputs, false).unwrap();
|
||||
|
||||
assert_eq!(results.len(), 3);
|
||||
|
||||
// With parallel execution, duplicate inputs may be processed simultaneously
|
||||
// and both see cache misses. Verify results are correct instead.
|
||||
assert_eq!(results[0].token_ids(), results[2].token_ids()); // Both "Hello" should match
|
||||
|
||||
// After batch processing, cache should be populated
|
||||
// Subsequent calls should hit the cache
|
||||
let _ = cached.encode("Hello", false).unwrap();
|
||||
let stats = cached.cache_stats().unwrap();
|
||||
|
||||
// Should have at least 1 hit from the call above (cache was populated by batch)
|
||||
assert!(
|
||||
stats.hits >= 1,
|
||||
"Expected at least 1 cache hit after batch processing"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_decoder_passthrough() {
|
||||
let tokenizer = Arc::new(MockTokenizer::new());
|
||||
let cached = CachedTokenizer::new(tokenizer, CacheConfig::default());
|
||||
|
||||
let tokens = vec![1, 2, 3];
|
||||
let decoded = cached.decode(&tokens, false).unwrap();
|
||||
|
||||
// Should just pass through to inner tokenizer
|
||||
assert!(!decoded.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_tokenizer_trait_methods() {
|
||||
let tokenizer = Arc::new(MockTokenizer::new());
|
||||
let cached = CachedTokenizer::new(tokenizer.clone(), CacheConfig::default());
|
||||
|
||||
// Should pass through to inner tokenizer
|
||||
assert_eq!(cached.vocab_size(), tokenizer.vocab_size());
|
||||
assert!(cached.token_to_id("Hello").is_some());
|
||||
assert!(cached.id_to_token(1).is_some());
|
||||
}
|
||||
}
|
||||
@@ -1,534 +0,0 @@
|
||||
//! Chat template support for tokenizers using Jinja2 templates
|
||||
//!
|
||||
//! This module provides functionality to apply chat templates to messages,
|
||||
//! similar to HuggingFace transformers' apply_chat_template method.
|
||||
|
||||
use std::{collections::HashMap, fs};
|
||||
|
||||
use anyhow::{anyhow, Result};
|
||||
use minijinja::{
|
||||
context,
|
||||
machinery::{
|
||||
ast::{Expr, Stmt},
|
||||
parse, WhitespaceConfig,
|
||||
},
|
||||
syntax::SyntaxConfig,
|
||||
value::Kwargs,
|
||||
Environment, Error as MinijinjaError, ErrorKind, Value,
|
||||
};
|
||||
use serde::Serialize;
|
||||
use serde_json::{self, ser::PrettyFormatter, Value as JsonValue}; // codespell:ignore ser
|
||||
|
||||
/// Chat template content format
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
|
||||
pub enum ChatTemplateContentFormat {
|
||||
/// Content is a simple string
|
||||
#[default]
|
||||
String,
|
||||
/// Content is a list of structured parts (OpenAI format)
|
||||
OpenAI,
|
||||
}
|
||||
|
||||
impl std::fmt::Display for ChatTemplateContentFormat {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
Self::String => write!(f, "string"),
|
||||
Self::OpenAI => write!(f, "openai"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Detect the content format expected by a Jinja2 chat template
|
||||
///
|
||||
/// This implements the same detection logic as SGLang's detect_jinja_template_content_format
|
||||
/// which uses AST parsing to look for content iteration patterns.
|
||||
///
|
||||
/// Returns:
|
||||
/// - ChatTemplateContentFormat::OpenAI if template expects structured content (list of parts)
|
||||
/// - ChatTemplateContentFormat::String if template expects simple string content
|
||||
pub fn detect_chat_template_content_format(template: &str) -> ChatTemplateContentFormat {
|
||||
// Use AST-based detection (enabled by default)
|
||||
if let Some(format) = detect_format_with_ast(template) {
|
||||
return format;
|
||||
}
|
||||
|
||||
// Default to string format if AST parsing fails
|
||||
ChatTemplateContentFormat::String
|
||||
}
|
||||
|
||||
/// Flags tracking which OpenAI-style patterns we've seen
|
||||
#[derive(Default, Debug, Clone, Copy)]
|
||||
struct Flags {
|
||||
saw_iteration: bool,
|
||||
saw_structure: bool,
|
||||
saw_assignment: bool,
|
||||
saw_macro: bool,
|
||||
}
|
||||
|
||||
impl Flags {
|
||||
fn any(self) -> bool {
|
||||
self.saw_iteration || self.saw_structure || self.saw_assignment || self.saw_macro
|
||||
}
|
||||
}
|
||||
|
||||
/// Single-pass AST detector with scope tracking
|
||||
struct Detector<'a> {
|
||||
ast: &'a Stmt<'a>,
|
||||
/// Message loop vars currently in scope (e.g., `message`, `m`, `msg`)
|
||||
scope: std::collections::VecDeque<String>,
|
||||
scope_set: std::collections::HashSet<String>,
|
||||
flags: Flags,
|
||||
}
|
||||
|
||||
impl<'a> Detector<'a> {
|
||||
fn new(ast: &'a Stmt<'a>) -> Self {
|
||||
Self {
|
||||
ast,
|
||||
scope: std::collections::VecDeque::new(),
|
||||
scope_set: std::collections::HashSet::new(),
|
||||
flags: Flags::default(),
|
||||
}
|
||||
}
|
||||
|
||||
fn run(mut self) -> Flags {
|
||||
self.walk_stmt(self.ast);
|
||||
self.flags
|
||||
}
|
||||
|
||||
fn push_scope(&mut self, var: String) {
|
||||
self.scope.push_back(var.clone());
|
||||
self.scope_set.insert(var);
|
||||
}
|
||||
|
||||
fn pop_scope(&mut self) {
|
||||
if let Some(v) = self.scope.pop_back() {
|
||||
self.scope_set.remove(&v);
|
||||
}
|
||||
}
|
||||
|
||||
fn is_var_access(expr: &Expr, varname: &str) -> bool {
|
||||
matches!(expr, Expr::Var(v) if v.id == varname)
|
||||
}
|
||||
|
||||
fn is_const_str(expr: &Expr, value: &str) -> bool {
|
||||
matches!(expr, Expr::Const(c) if c.value.as_str() == Some(value))
|
||||
}
|
||||
|
||||
fn is_numeric_const(expr: &Expr) -> bool {
|
||||
matches!(expr, Expr::Const(c) if c.value.is_number())
|
||||
}
|
||||
|
||||
/// Check if expr is varname.content or varname["content"]
|
||||
fn is_var_dot_content(expr: &Expr, varname: &str) -> bool {
|
||||
match expr {
|
||||
Expr::GetAttr(g) => Self::is_var_access(&g.expr, varname) && g.name == "content",
|
||||
Expr::GetItem(g) => {
|
||||
Self::is_var_access(&g.expr, varname)
|
||||
&& Self::is_const_str(&g.subscript_expr, "content")
|
||||
}
|
||||
// Unwrap filters/tests that just wrap the same expr
|
||||
Expr::Filter(f) => f
|
||||
.expr
|
||||
.as_ref()
|
||||
.is_some_and(|e| Self::is_var_dot_content(e, varname)),
|
||||
Expr::Test(t) => Self::is_var_dot_content(&t.expr, varname),
|
||||
_ => false,
|
||||
}
|
||||
}
|
||||
|
||||
/// Check if expr accesses .content on any variable in our scope, or any descendant of it.
|
||||
fn is_any_scope_var_content(&self, expr: &Expr) -> bool {
|
||||
let mut current_expr = expr;
|
||||
loop {
|
||||
// Check if current level matches <scopeVar>.content
|
||||
if self
|
||||
.scope_set
|
||||
.iter()
|
||||
.any(|v| Self::is_var_dot_content(current_expr, v))
|
||||
{
|
||||
return true;
|
||||
}
|
||||
// Walk up the expression tree
|
||||
match current_expr {
|
||||
Expr::GetAttr(g) => current_expr = &g.expr,
|
||||
Expr::GetItem(g) => current_expr = &g.expr,
|
||||
_ => return false,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn walk_stmt(&mut self, stmt: &Stmt) {
|
||||
// Early exit if we've already detected an OpenAI pattern
|
||||
if self.flags.any() {
|
||||
return;
|
||||
}
|
||||
|
||||
match stmt {
|
||||
Stmt::Template(t) => {
|
||||
for ch in &t.children {
|
||||
self.walk_stmt(ch);
|
||||
}
|
||||
}
|
||||
// {% for message in messages %}
|
||||
Stmt::ForLoop(fl) => {
|
||||
// Detect "for X in messages" → push X into scope
|
||||
if let Expr::Var(iter) = &fl.iter {
|
||||
if iter.id == "messages" {
|
||||
if let Expr::Var(target) = &fl.target {
|
||||
self.push_scope(target.id.to_string());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Also detect "for ... in message.content" or "for ... in content"
|
||||
// - Iterating directly over <scopeVar>.content => OpenAI style
|
||||
if self.is_any_scope_var_content(&fl.iter) {
|
||||
self.flags.saw_iteration = true;
|
||||
}
|
||||
// - Iterating over a local var named "content"
|
||||
if matches!(&fl.iter, Expr::Var(v) if v.id == "content") {
|
||||
self.flags.saw_iteration = true;
|
||||
}
|
||||
|
||||
for b in &fl.body {
|
||||
self.walk_stmt(b);
|
||||
}
|
||||
|
||||
// Pop scope if we pushed it
|
||||
if let Expr::Var(iter) = &fl.iter {
|
||||
if iter.id == "messages" && matches!(&fl.target, Expr::Var(_)) {
|
||||
self.pop_scope();
|
||||
}
|
||||
}
|
||||
}
|
||||
Stmt::IfCond(ic) => {
|
||||
self.inspect_expr_for_structure(&ic.expr);
|
||||
for b in &ic.true_body {
|
||||
self.walk_stmt(b);
|
||||
}
|
||||
for b in &ic.false_body {
|
||||
self.walk_stmt(b);
|
||||
}
|
||||
}
|
||||
Stmt::EmitExpr(e) => {
|
||||
self.inspect_expr_for_structure(&e.expr);
|
||||
}
|
||||
// {% set content = message.content %}
|
||||
Stmt::Set(s) => {
|
||||
if Self::is_var_access(&s.target, "content")
|
||||
&& self.is_any_scope_var_content(&s.expr)
|
||||
{
|
||||
self.flags.saw_assignment = true;
|
||||
}
|
||||
}
|
||||
Stmt::Macro(m) => {
|
||||
// Heuristic: macro that checks type (via `is` test) and also has any loop
|
||||
let mut has_type_check = false;
|
||||
let mut has_loop = false;
|
||||
Self::scan_macro_body(&m.body, &mut has_type_check, &mut has_loop);
|
||||
if has_type_check && has_loop {
|
||||
self.flags.saw_macro = true;
|
||||
}
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
fn inspect_expr_for_structure(&mut self, expr: &Expr) {
|
||||
if self.flags.saw_structure {
|
||||
return;
|
||||
}
|
||||
|
||||
match expr {
|
||||
// content[0] or message.content[0]
|
||||
Expr::GetItem(gi) => {
|
||||
if (matches!(&gi.expr, Expr::Var(v) if v.id == "content")
|
||||
|| self.is_any_scope_var_content(&gi.expr))
|
||||
&& Self::is_numeric_const(&gi.subscript_expr)
|
||||
{
|
||||
self.flags.saw_structure = true;
|
||||
}
|
||||
}
|
||||
// content|length or message.content|length
|
||||
Expr::Filter(f) => {
|
||||
if f.name == "length" {
|
||||
if let Some(inner) = &f.expr {
|
||||
// Box derefs automatically, so `&**inner` is `&Expr`
|
||||
let inner_ref: &Expr = inner;
|
||||
let is_content_var = matches!(inner_ref, Expr::Var(v) if v.id == "content");
|
||||
if is_content_var || self.is_any_scope_var_content(inner_ref) {
|
||||
self.flags.saw_structure = true;
|
||||
}
|
||||
}
|
||||
} else if let Some(inner) = &f.expr {
|
||||
let inner_ref: &Expr = inner;
|
||||
self.inspect_expr_for_structure(inner_ref);
|
||||
}
|
||||
}
|
||||
// content is sequence/iterable OR message.content is sequence/iterable
|
||||
Expr::Test(t) => {
|
||||
if t.name == "sequence" || t.name == "iterable" || t.name == "string" {
|
||||
if matches!(&t.expr, Expr::Var(v) if v.id == "content")
|
||||
|| self.is_any_scope_var_content(&t.expr)
|
||||
{
|
||||
self.flags.saw_structure = true;
|
||||
}
|
||||
} else {
|
||||
self.inspect_expr_for_structure(&t.expr);
|
||||
}
|
||||
}
|
||||
Expr::GetAttr(g) => {
|
||||
// Keep walking; nested expressions can hide structure checks
|
||||
self.inspect_expr_for_structure(&g.expr);
|
||||
}
|
||||
// Handle binary operations like: if (message.content is string) and other_cond
|
||||
Expr::BinOp(op) => {
|
||||
self.inspect_expr_for_structure(&op.left);
|
||||
self.inspect_expr_for_structure(&op.right);
|
||||
}
|
||||
// Handle unary operations like: if not (message.content is string)
|
||||
Expr::UnaryOp(op) => {
|
||||
self.inspect_expr_for_structure(&op.expr);
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
fn scan_macro_body(body: &[Stmt], has_type_check: &mut bool, has_loop: &mut bool) {
|
||||
for s in body {
|
||||
if *has_type_check && *has_loop {
|
||||
return;
|
||||
}
|
||||
|
||||
match s {
|
||||
Stmt::IfCond(ic) => {
|
||||
if matches!(&ic.expr, Expr::Test(_)) {
|
||||
*has_type_check = true;
|
||||
}
|
||||
Self::scan_macro_body(&ic.true_body, has_type_check, has_loop);
|
||||
Self::scan_macro_body(&ic.false_body, has_type_check, has_loop);
|
||||
}
|
||||
Stmt::ForLoop(fl) => {
|
||||
*has_loop = true;
|
||||
Self::scan_macro_body(&fl.body, has_type_check, has_loop);
|
||||
}
|
||||
Stmt::Template(t) => {
|
||||
Self::scan_macro_body(&t.children, has_type_check, has_loop);
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// AST-based detection using minijinja's unstable machinery
|
||||
/// Single-pass detector with scope tracking
|
||||
fn detect_format_with_ast(template: &str) -> Option<ChatTemplateContentFormat> {
|
||||
let ast = match parse(
|
||||
template,
|
||||
"template",
|
||||
SyntaxConfig {},
|
||||
WhitespaceConfig::default(),
|
||||
) {
|
||||
Ok(ast) => ast,
|
||||
Err(_) => return Some(ChatTemplateContentFormat::String),
|
||||
};
|
||||
|
||||
let flags = Detector::new(&ast).run();
|
||||
Some(if flags.any() {
|
||||
ChatTemplateContentFormat::OpenAI
|
||||
} else {
|
||||
ChatTemplateContentFormat::String
|
||||
})
|
||||
}
|
||||
|
||||
/// Parameters for chat template application
|
||||
#[derive(Default)]
|
||||
pub struct ChatTemplateParams<'a> {
|
||||
pub add_generation_prompt: bool,
|
||||
pub tools: Option<&'a [serde_json::Value]>,
|
||||
pub documents: Option<&'a [serde_json::Value]>,
|
||||
pub template_kwargs: Option<&'a HashMap<String, serde_json::Value>>,
|
||||
}
|
||||
|
||||
/// Custom tojson filter compatible with HuggingFace transformers' implementation.
|
||||
///
|
||||
/// HuggingFace transformers registers a custom `tojson` filter that accepts additional
|
||||
/// keyword arguments beyond what standard Jinja2 provides:
|
||||
/// - `ensure_ascii` (bool): Whether to escape non-ASCII characters (ignored in Rust, always UTF-8)
|
||||
/// - `indent` (int): Number of spaces for indentation (pretty-printing)
|
||||
/// - `separators` (ignored): Custom separators for JSON output
|
||||
/// - `sort_keys` (bool): Whether to sort dictionary keys
|
||||
///
|
||||
/// This is necessary for compatibility with chat templates from HuggingFace Hub models.
|
||||
/// See: https://github.com/huggingface/transformers/blob/main/src/transformers/utils/chat_template_utils.py
|
||||
fn tojson_filter(value: Value, kwargs: Kwargs) -> std::result::Result<Value, MinijinjaError> {
|
||||
let _ensure_ascii: Option<bool> = kwargs.get("ensure_ascii")?;
|
||||
let indent: Option<i64> = kwargs.get("indent")?;
|
||||
let _separators: Option<Value> = kwargs.get("separators")?;
|
||||
let sort_keys: Option<bool> = kwargs.get("sort_keys")?;
|
||||
|
||||
// Ensure all kwargs are consumed to avoid "unknown keyword argument" errors
|
||||
kwargs.assert_all_used()?;
|
||||
|
||||
let json_value: serde_json::Value = serde_json::to_value(&value).map_err(|e| {
|
||||
MinijinjaError::new(
|
||||
ErrorKind::InvalidOperation,
|
||||
format!("Failed to convert to JSON value: {}", e),
|
||||
)
|
||||
})?;
|
||||
|
||||
// Helper to serialize with custom indentation
|
||||
fn serialize_with_indent<T: Serialize>(
|
||||
value: &T,
|
||||
spaces: usize,
|
||||
) -> std::result::Result<String, MinijinjaError> {
|
||||
let indent_str = vec![b' '; spaces];
|
||||
let formatter = PrettyFormatter::with_indent(&indent_str);
|
||||
let mut buf = Vec::new();
|
||||
let mut serializer = serde_json::Serializer::with_formatter(&mut buf, formatter);
|
||||
value.serialize(&mut serializer).map_err(|e| {
|
||||
MinijinjaError::new(
|
||||
ErrorKind::InvalidOperation,
|
||||
format!("Failed to serialize JSON: {}", e),
|
||||
)
|
||||
})?;
|
||||
String::from_utf8(buf).map_err(|e| {
|
||||
MinijinjaError::new(
|
||||
ErrorKind::InvalidOperation,
|
||||
format!("Invalid UTF-8 in JSON output: {}", e),
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
// Serialize with options
|
||||
let json_str: std::result::Result<String, MinijinjaError> = {
|
||||
let sorted_json;
|
||||
let value_to_serialize = if sort_keys.unwrap_or(false) {
|
||||
sorted_json = sort_json_keys(&json_value);
|
||||
&sorted_json
|
||||
} else {
|
||||
&json_value
|
||||
};
|
||||
|
||||
if let Some(spaces) = indent {
|
||||
if spaces < 0 {
|
||||
return Err(MinijinjaError::new(
|
||||
ErrorKind::InvalidOperation,
|
||||
"indent cannot be negative",
|
||||
));
|
||||
}
|
||||
serialize_with_indent(value_to_serialize, spaces as usize)
|
||||
} else {
|
||||
serde_json::to_string(value_to_serialize).map_err(|e| {
|
||||
MinijinjaError::new(
|
||||
ErrorKind::InvalidOperation,
|
||||
format!("Failed to serialize JSON: {}", e),
|
||||
)
|
||||
})
|
||||
}
|
||||
};
|
||||
|
||||
json_str.map(Value::from_safe_string)
|
||||
}
|
||||
|
||||
/// Recursively sort all object keys in a JSON value
|
||||
fn sort_json_keys(value: &JsonValue) -> JsonValue {
|
||||
match value {
|
||||
JsonValue::Object(map) => {
|
||||
let mut sorted: serde_json::Map<String, JsonValue> = serde_json::Map::new();
|
||||
let mut keys: Vec<_> = map.keys().collect();
|
||||
keys.sort();
|
||||
for key in keys {
|
||||
sorted.insert(key.clone(), sort_json_keys(&map[key]));
|
||||
}
|
||||
JsonValue::Object(sorted)
|
||||
}
|
||||
JsonValue::Array(arr) => JsonValue::Array(arr.iter().map(sort_json_keys).collect()),
|
||||
_ => value.clone(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Chat template processor using Jinja2 - simple wrapper like HuggingFace
|
||||
pub struct ChatTemplateProcessor {
|
||||
template: String,
|
||||
}
|
||||
|
||||
impl ChatTemplateProcessor {
|
||||
/// Create a new chat template processor
|
||||
pub fn new(template: String) -> Self {
|
||||
ChatTemplateProcessor { template }
|
||||
}
|
||||
|
||||
/// Apply the chat template to a list of messages
|
||||
///
|
||||
/// This mimics the behavior of HuggingFace's apply_chat_template method
|
||||
/// but returns the formatted string instead of token IDs.
|
||||
/// Messages should be pre-processed into the format expected by the template.
|
||||
pub fn apply_chat_template(
|
||||
&self,
|
||||
messages: &[serde_json::Value],
|
||||
params: ChatTemplateParams,
|
||||
) -> Result<String> {
|
||||
let mut env = Environment::new();
|
||||
|
||||
// Register the template
|
||||
env.add_template("chat", &self.template)
|
||||
.map_err(|e| anyhow!("Failed to add template: {}", e))?;
|
||||
|
||||
// Enable Python method compatibility (e.g., str.startswith, str.endswith)
|
||||
env.set_unknown_method_callback(minijinja_contrib::pycompat::unknown_method_callback);
|
||||
|
||||
// Register custom tojson filter compatible with HuggingFace transformers
|
||||
// This overrides minijinja's built-in tojson to support additional kwargs
|
||||
// like ensure_ascii, separators, and sort_keys that HuggingFace templates use
|
||||
env.add_filter("tojson", tojson_filter);
|
||||
|
||||
// Get the template
|
||||
let tmpl = env
|
||||
.get_template("chat")
|
||||
.map_err(|e| anyhow!("Failed to get template: {}", e))?;
|
||||
|
||||
// Convert messages to minijinja::Value (messages already processed by router)
|
||||
let minijinja_messages: Vec<Value> = messages.iter().map(Value::from_serialize).collect();
|
||||
|
||||
let base_context = context! {
|
||||
messages => &minijinja_messages,
|
||||
add_generation_prompt => params.add_generation_prompt,
|
||||
tools => params.tools,
|
||||
documents => params.documents,
|
||||
};
|
||||
|
||||
// Merge with template_kwargs if provided
|
||||
let ctx = if let Some(kwargs) = params.template_kwargs {
|
||||
context! {
|
||||
..base_context,
|
||||
..Value::from_serialize(kwargs)
|
||||
}
|
||||
} else {
|
||||
base_context
|
||||
};
|
||||
|
||||
// Render the template
|
||||
let rendered = tmpl
|
||||
.render(&ctx)
|
||||
.map_err(|e| anyhow!("Failed to render template: {}", e))?;
|
||||
|
||||
Ok(rendered)
|
||||
}
|
||||
}
|
||||
|
||||
/// Load chat template from tokenizer config JSON
|
||||
pub fn load_chat_template_from_config(config_path: &str) -> Result<Option<String>> {
|
||||
let content = fs::read_to_string(config_path)?;
|
||||
let config: serde_json::Value = serde_json::from_str(&content)?;
|
||||
|
||||
// Look for chat_template in the config
|
||||
if let Some(template) = config.get("chat_template") {
|
||||
if let Some(template_str) = template.as_str() {
|
||||
return Ok(Some(template_str.to_string()));
|
||||
}
|
||||
}
|
||||
|
||||
Ok(None)
|
||||
}
|
||||
@@ -1,484 +0,0 @@
|
||||
use std::{fs::File, io::Read, path::Path, sync::Arc};
|
||||
|
||||
use anyhow::{Error, Result};
|
||||
use tracing::debug;
|
||||
|
||||
use super::{huggingface::HuggingFaceTokenizer, tiktoken::TiktokenTokenizer, traits};
|
||||
use crate::tokenizer::hub::download_tokenizer_from_hf;
|
||||
|
||||
/// Represents the type of tokenizer being used
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum TokenizerType {
|
||||
HuggingFace(String),
|
||||
Mock,
|
||||
Tiktoken(String),
|
||||
// Future: SentencePiece, GGUF
|
||||
}
|
||||
|
||||
/// Create a tokenizer from a file path to a tokenizer file.
|
||||
/// The file extension is used to determine the tokenizer type.
|
||||
/// Supported file types are:
|
||||
/// - json: HuggingFace tokenizer
|
||||
/// - For testing: can return mock tokenizer
|
||||
pub fn create_tokenizer_from_file(file_path: &str) -> Result<Arc<dyn traits::Tokenizer>> {
|
||||
create_tokenizer_with_chat_template(file_path, None)
|
||||
}
|
||||
|
||||
/// Create a tokenizer from a file path with an optional chat template
|
||||
pub fn create_tokenizer_with_chat_template(
|
||||
file_path: &str,
|
||||
chat_template_path: Option<&str>,
|
||||
) -> Result<Arc<dyn traits::Tokenizer>> {
|
||||
// Special case for testing
|
||||
if file_path == "mock" || file_path == "test" {
|
||||
return Ok(Arc::new(super::mock::MockTokenizer::new()));
|
||||
}
|
||||
|
||||
let path = Path::new(file_path);
|
||||
|
||||
// Check if file exists
|
||||
if !path.exists() {
|
||||
return Err(Error::msg(format!("File not found: {}", file_path)));
|
||||
}
|
||||
|
||||
// If path is a directory, search for tokenizer files
|
||||
if path.is_dir() {
|
||||
let tokenizer_json = path.join("tokenizer.json");
|
||||
if tokenizer_json.exists() {
|
||||
// Resolve chat template: provided path takes precedence over auto-discovery
|
||||
let final_chat_template =
|
||||
resolve_and_log_chat_template(chat_template_path, path, file_path);
|
||||
let tokenizer_path_str = tokenizer_json.to_str().ok_or_else(|| {
|
||||
Error::msg(format!(
|
||||
"Tokenizer path is not valid UTF-8: {:?}",
|
||||
tokenizer_json
|
||||
))
|
||||
})?;
|
||||
return create_tokenizer_with_chat_template(
|
||||
tokenizer_path_str,
|
||||
final_chat_template.as_deref(),
|
||||
);
|
||||
}
|
||||
|
||||
return Err(Error::msg(format!(
|
||||
"Directory '{}' does not contain a valid tokenizer file (tokenizer.json, tokenizer_config.json, or vocab.json)",
|
||||
file_path
|
||||
)));
|
||||
}
|
||||
|
||||
// Try to determine tokenizer type from extension
|
||||
let extension = path
|
||||
.extension()
|
||||
.and_then(std::ffi::OsStr::to_str)
|
||||
.map(|s| s.to_lowercase());
|
||||
|
||||
let result = match extension.as_deref() {
|
||||
Some("json") => {
|
||||
let tokenizer =
|
||||
HuggingFaceTokenizer::from_file_with_chat_template(file_path, chat_template_path)?;
|
||||
|
||||
Ok(Arc::new(tokenizer) as Arc<dyn traits::Tokenizer>)
|
||||
}
|
||||
Some("model") => {
|
||||
// SentencePiece model file
|
||||
Err(Error::msg("SentencePiece models not yet supported"))
|
||||
}
|
||||
Some("gguf") => {
|
||||
// GGUF format
|
||||
Err(Error::msg("GGUF format not yet supported"))
|
||||
}
|
||||
_ => {
|
||||
// Try to auto-detect by reading file content
|
||||
auto_detect_tokenizer(file_path)
|
||||
}
|
||||
};
|
||||
|
||||
result
|
||||
}
|
||||
|
||||
/// Auto-detect tokenizer type by examining file content
|
||||
fn auto_detect_tokenizer(file_path: &str) -> Result<Arc<dyn traits::Tokenizer>> {
|
||||
let mut file = File::open(file_path)?;
|
||||
let mut buffer = vec![0u8; 512]; // Read first 512 bytes for detection
|
||||
let bytes_read = file.read(&mut buffer)?;
|
||||
buffer.truncate(bytes_read);
|
||||
|
||||
// Check for JSON (HuggingFace format)
|
||||
if is_likely_json(&buffer) {
|
||||
let tokenizer = HuggingFaceTokenizer::from_file(file_path)?;
|
||||
return Ok(Arc::new(tokenizer));
|
||||
}
|
||||
|
||||
// Check for GGUF magic number
|
||||
if buffer.len() >= 4 && &buffer[0..4] == b"GGUF" {
|
||||
return Err(Error::msg("GGUF format detected but not yet supported"));
|
||||
}
|
||||
|
||||
// Check for SentencePiece model
|
||||
if is_likely_sentencepiece(&buffer) {
|
||||
return Err(Error::msg(
|
||||
"SentencePiece model detected but not yet supported",
|
||||
));
|
||||
}
|
||||
|
||||
Err(Error::msg(format!(
|
||||
"Unable to determine tokenizer type for file: {}",
|
||||
file_path
|
||||
)))
|
||||
}
|
||||
|
||||
/// Check if the buffer likely contains JSON data
|
||||
fn is_likely_json(buffer: &[u8]) -> bool {
|
||||
// Skip UTF-8 BOM if present
|
||||
let content = if buffer.len() >= 3 && buffer[0..3] == [0xEF, 0xBB, 0xBF] {
|
||||
&buffer[3..]
|
||||
} else {
|
||||
buffer
|
||||
};
|
||||
|
||||
// Find first non-whitespace character without allocation
|
||||
if let Some(first_byte) = content.iter().find(|&&b| !b.is_ascii_whitespace()) {
|
||||
*first_byte == b'{' || *first_byte == b'['
|
||||
} else {
|
||||
false
|
||||
}
|
||||
}
|
||||
|
||||
/// Check if the buffer likely contains a SentencePiece model
|
||||
fn is_likely_sentencepiece(buffer: &[u8]) -> bool {
|
||||
// SentencePiece models often start with specific patterns
|
||||
// This is a simplified check
|
||||
if buffer.len() < 12 {
|
||||
return false;
|
||||
}
|
||||
|
||||
// Check header patterns first (cheap)
|
||||
if buffer.starts_with(b"\x0a\x09") || buffer.starts_with(b"\x08\x00") {
|
||||
return true;
|
||||
}
|
||||
|
||||
// Single-pass scan for special token markers
|
||||
// Instead of multiple windows() calls, scan once looking for all patterns
|
||||
let patterns: &[&[u8]] = &[b"<unk", b"<s>", b"</s>"];
|
||||
for window in buffer.windows(4) {
|
||||
for pattern in patterns {
|
||||
if window.starts_with(pattern) {
|
||||
return true;
|
||||
}
|
||||
}
|
||||
}
|
||||
false
|
||||
}
|
||||
|
||||
/// Helper function to discover chat template files in a directory
|
||||
pub fn discover_chat_template_in_dir(dir: &Path) -> Option<String> {
|
||||
use std::fs;
|
||||
|
||||
// Priority 1: Look for chat_template.json (contains Jinja in JSON format)
|
||||
let json_template_path = dir.join("chat_template.json");
|
||||
if json_template_path.exists() {
|
||||
return json_template_path.to_str().map(|s| s.to_string());
|
||||
}
|
||||
|
||||
// Priority 2: Look for chat_template.jinja (standard Jinja file)
|
||||
let jinja_path = dir.join("chat_template.jinja");
|
||||
if jinja_path.exists() {
|
||||
return jinja_path.to_str().map(|s| s.to_string());
|
||||
}
|
||||
|
||||
// Priority 3: Look for any .jinja file (for models with non-standard naming)
|
||||
if let Ok(entries) = fs::read_dir(dir) {
|
||||
for entry in entries.flatten() {
|
||||
if let Some(name) = entry.file_name().to_str() {
|
||||
if name.ends_with(".jinja") && name != "chat_template.jinja" {
|
||||
return entry.path().to_str().map(|s| s.to_string());
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
None
|
||||
}
|
||||
|
||||
/// Helper function to resolve and log chat template selection
|
||||
///
|
||||
/// Resolves the final chat template to use by prioritizing provided path over auto-discovery,
|
||||
/// and logs the source for debugging purposes.
|
||||
fn resolve_and_log_chat_template(
|
||||
provided_path: Option<&str>,
|
||||
discovery_dir: &Path,
|
||||
model_name: &str,
|
||||
) -> Option<String> {
|
||||
let final_chat_template = provided_path
|
||||
.map(|s| s.to_string())
|
||||
.or_else(|| discover_chat_template_in_dir(discovery_dir));
|
||||
|
||||
match (&provided_path, &final_chat_template) {
|
||||
(Some(provided), _) => {
|
||||
debug!("Using provided chat template: {}", provided);
|
||||
}
|
||||
(None, Some(discovered)) => {
|
||||
debug!(
|
||||
"Auto-discovered chat template in '{}': {}",
|
||||
discovery_dir.display(),
|
||||
discovered
|
||||
);
|
||||
}
|
||||
(None, None) => {
|
||||
debug!(
|
||||
"No chat template provided or discovered for model: {}",
|
||||
model_name
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
final_chat_template
|
||||
}
|
||||
|
||||
/// Factory function to create tokenizer from a model name or path (async version)
|
||||
pub async fn create_tokenizer_async(
|
||||
model_name_or_path: &str,
|
||||
) -> Result<Arc<dyn traits::Tokenizer>> {
|
||||
create_tokenizer_async_with_chat_template(model_name_or_path, None).await
|
||||
}
|
||||
|
||||
/// Factory function to create tokenizer with optional chat template (async version)
|
||||
pub async fn create_tokenizer_async_with_chat_template(
|
||||
model_name_or_path: &str,
|
||||
chat_template_path: Option<&str>,
|
||||
) -> Result<Arc<dyn traits::Tokenizer>> {
|
||||
// Check if it's a file path
|
||||
let path = Path::new(model_name_or_path);
|
||||
if path.exists() {
|
||||
return create_tokenizer_with_chat_template(model_name_or_path, chat_template_path);
|
||||
}
|
||||
|
||||
// Check if it's a GPT model name that should use Tiktoken
|
||||
// Only match specific OpenAI model patterns to avoid catching HuggingFace models like "openai/gpt-oss-20b"
|
||||
if model_name_or_path.contains("gpt-4")
|
||||
|| model_name_or_path.contains("gpt-3.5")
|
||||
|| model_name_or_path.contains("gpt-3")
|
||||
|| model_name_or_path.contains("turbo")
|
||||
|| model_name_or_path.contains("davinci")
|
||||
|| model_name_or_path.contains("curie")
|
||||
|| model_name_or_path.contains("babbage")
|
||||
|| model_name_or_path.contains("ada")
|
||||
|| model_name_or_path.contains("codex")
|
||||
{
|
||||
// Try tiktoken first, but fall back to HuggingFace if it fails
|
||||
match TiktokenTokenizer::from_model_name(model_name_or_path) {
|
||||
Ok(tokenizer) => return Ok(Arc::new(tokenizer)),
|
||||
Err(e) => {
|
||||
debug!(
|
||||
"Tiktoken failed for '{}': {}, falling back to HuggingFace",
|
||||
model_name_or_path, e
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Try to download tokenizer files from HuggingFace
|
||||
match download_tokenizer_from_hf(model_name_or_path).await {
|
||||
Ok(cache_dir) => {
|
||||
// Look for tokenizer.json in the cache directory
|
||||
let tokenizer_path = cache_dir.join("tokenizer.json");
|
||||
if tokenizer_path.exists() {
|
||||
// Resolve chat template: provided path takes precedence over auto-discovery
|
||||
let final_chat_template = resolve_and_log_chat_template(
|
||||
chat_template_path,
|
||||
&cache_dir,
|
||||
model_name_or_path,
|
||||
);
|
||||
|
||||
let tokenizer_path_str = tokenizer_path.to_str().ok_or_else(|| {
|
||||
Error::msg(format!(
|
||||
"Tokenizer path is not valid UTF-8: {:?}",
|
||||
tokenizer_path
|
||||
))
|
||||
})?;
|
||||
create_tokenizer_with_chat_template(
|
||||
tokenizer_path_str,
|
||||
final_chat_template.as_deref(),
|
||||
)
|
||||
} else {
|
||||
// Try other common tokenizer file names
|
||||
let possible_files = ["tokenizer_config.json", "vocab.json"];
|
||||
for file_name in &possible_files {
|
||||
let file_path = cache_dir.join(file_name);
|
||||
if file_path.exists() {
|
||||
// Resolve chat template: provided path takes precedence over auto-discovery
|
||||
let final_chat_template = resolve_and_log_chat_template(
|
||||
chat_template_path,
|
||||
&cache_dir,
|
||||
model_name_or_path,
|
||||
);
|
||||
|
||||
let file_path_str = file_path.to_str().ok_or_else(|| {
|
||||
Error::msg(format!("File path is not valid UTF-8: {:?}", file_path))
|
||||
})?;
|
||||
return create_tokenizer_with_chat_template(
|
||||
file_path_str,
|
||||
final_chat_template.as_deref(),
|
||||
);
|
||||
}
|
||||
}
|
||||
Err(Error::msg(format!(
|
||||
"Downloaded model '{}' but couldn't find a suitable tokenizer file",
|
||||
model_name_or_path
|
||||
)))
|
||||
}
|
||||
}
|
||||
Err(e) => Err(Error::msg(format!(
|
||||
"Failed to download tokenizer from HuggingFace: {}",
|
||||
e
|
||||
))),
|
||||
}
|
||||
}
|
||||
|
||||
/// Factory function to create tokenizer from a model name or path (blocking version)
|
||||
///
|
||||
/// This delegates to `create_tokenizer_with_chat_template_blocking` with no chat template,
|
||||
/// which handles both local files and HuggingFace Hub downloads uniformly.
|
||||
pub fn create_tokenizer(model_name_or_path: &str) -> Result<Arc<dyn traits::Tokenizer>> {
|
||||
create_tokenizer_with_chat_template_blocking(model_name_or_path, None)
|
||||
}
|
||||
|
||||
/// Factory function to create tokenizer with optional chat template (blocking version)
|
||||
pub fn create_tokenizer_with_chat_template_blocking(
|
||||
model_name_or_path: &str,
|
||||
chat_template_path: Option<&str>,
|
||||
) -> Result<Arc<dyn traits::Tokenizer>> {
|
||||
// Check if it's a file path
|
||||
let path = Path::new(model_name_or_path);
|
||||
if path.exists() {
|
||||
return create_tokenizer_with_chat_template(model_name_or_path, chat_template_path);
|
||||
}
|
||||
|
||||
// Check if it's a GPT model name that should use Tiktoken
|
||||
if model_name_or_path.contains("gpt-")
|
||||
|| model_name_or_path.contains("davinci")
|
||||
|| model_name_or_path.contains("curie")
|
||||
|| model_name_or_path.contains("babbage")
|
||||
|| model_name_or_path.contains("ada")
|
||||
{
|
||||
let tokenizer = TiktokenTokenizer::from_model_name(model_name_or_path)?;
|
||||
return Ok(Arc::new(tokenizer));
|
||||
}
|
||||
|
||||
// Only use tokio for HuggingFace downloads
|
||||
// Check if we're already in a tokio runtime
|
||||
if let Ok(handle) = tokio::runtime::Handle::try_current() {
|
||||
// We're in a runtime, use block_in_place
|
||||
tokio::task::block_in_place(|| {
|
||||
handle.block_on(create_tokenizer_async_with_chat_template(
|
||||
model_name_or_path,
|
||||
chat_template_path,
|
||||
))
|
||||
})
|
||||
} else {
|
||||
// No runtime, create a temporary one
|
||||
let rt = tokio::runtime::Runtime::new()?;
|
||||
rt.block_on(create_tokenizer_async_with_chat_template(
|
||||
model_name_or_path,
|
||||
chat_template_path,
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
/// Get information about a tokenizer file
|
||||
pub fn get_tokenizer_info(file_path: &str) -> Result<TokenizerType> {
|
||||
let path = Path::new(file_path);
|
||||
|
||||
if !path.exists() {
|
||||
return Err(Error::msg(format!("File not found: {}", file_path)));
|
||||
}
|
||||
|
||||
let extension = path
|
||||
.extension()
|
||||
.and_then(std::ffi::OsStr::to_str)
|
||||
.map(|s| s.to_lowercase());
|
||||
|
||||
match extension.as_deref() {
|
||||
Some("json") => Ok(TokenizerType::HuggingFace(file_path.to_string())),
|
||||
_ => {
|
||||
// Try auto-detection
|
||||
use std::{fs::File, io::Read};
|
||||
|
||||
let mut file = File::open(file_path)?;
|
||||
let mut buffer = vec![0u8; 512];
|
||||
let bytes_read = file.read(&mut buffer)?;
|
||||
buffer.truncate(bytes_read);
|
||||
|
||||
if is_likely_json(&buffer) {
|
||||
Ok(TokenizerType::HuggingFace(file_path.to_string()))
|
||||
} else {
|
||||
Err(Error::msg("Unknown tokenizer type"))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_json_detection() {
|
||||
assert!(is_likely_json(b"{\"test\": \"value\"}"));
|
||||
assert!(is_likely_json(b" \n\t{\"test\": \"value\"}"));
|
||||
assert!(is_likely_json(b"[1, 2, 3]"));
|
||||
assert!(!is_likely_json(b"not json"));
|
||||
assert!(!is_likely_json(b""));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_mock_tokenizer_creation() {
|
||||
let tokenizer = create_tokenizer_from_file("mock").unwrap();
|
||||
assert_eq!(tokenizer.vocab_size(), 14); // Mock tokenizer has 14 tokens
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_file_not_found() {
|
||||
let result = create_tokenizer_from_file("/nonexistent/file.json");
|
||||
assert!(result.is_err());
|
||||
if let Err(e) = result {
|
||||
assert!(e.to_string().contains("File not found"));
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_create_tiktoken_tokenizer() {
|
||||
let tokenizer = create_tokenizer("gpt-4").unwrap();
|
||||
assert!(tokenizer.vocab_size() > 0);
|
||||
|
||||
let text = "Hello, world!";
|
||||
let encoding = tokenizer.encode(text, false).unwrap();
|
||||
let decoded = tokenizer.decode(encoding.token_ids(), false).unwrap();
|
||||
assert_eq!(decoded, text);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_download_tokenizer_from_hf() {
|
||||
// Skip this test if HF_TOKEN is not set and we're in CI
|
||||
if std::env::var("CI").is_ok() && std::env::var("HF_TOKEN").is_err() {
|
||||
println!("Skipping HF download test in CI without HF_TOKEN");
|
||||
return;
|
||||
}
|
||||
|
||||
// Try to create tokenizer for a known small model
|
||||
let result = create_tokenizer_async("bert-base-uncased").await;
|
||||
|
||||
// The test might fail due to network issues or rate limiting
|
||||
// so we just check that the function executes without panic
|
||||
match result {
|
||||
Ok(tokenizer) => {
|
||||
assert!(tokenizer.vocab_size() > 0);
|
||||
println!("Successfully downloaded and created tokenizer");
|
||||
}
|
||||
Err(e) => {
|
||||
println!("Download failed (this might be expected): {}", e);
|
||||
// Don't fail the test - network issues shouldn't break CI
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,320 +0,0 @@
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
use hf_hub::api::tokio::ApiBuilder;
|
||||
|
||||
const IGNORED: [&str; 5] = [
|
||||
".gitattributes",
|
||||
"LICENSE",
|
||||
"LICENSE.txt",
|
||||
"README.md",
|
||||
"USE_POLICY.md",
|
||||
];
|
||||
|
||||
/// Checks if a file is a model weight file
|
||||
fn is_weight_file(filename: &str) -> bool {
|
||||
filename.ends_with(".bin")
|
||||
|| filename.ends_with(".safetensors")
|
||||
|| filename.ends_with(".h5")
|
||||
|| filename.ends_with(".msgpack")
|
||||
|| filename.ends_with(".ckpt.index")
|
||||
}
|
||||
|
||||
/// Checks if a file is an image file
|
||||
fn is_image(filename: &str) -> bool {
|
||||
filename.ends_with(".png")
|
||||
|| filename.ends_with("PNG")
|
||||
|| filename.ends_with(".jpg")
|
||||
|| filename.ends_with("JPG")
|
||||
|| filename.ends_with(".jpeg")
|
||||
|| filename.ends_with("JPEG")
|
||||
}
|
||||
|
||||
/// Checks if a file is a tokenizer file
|
||||
fn is_tokenizer_file(filename: &str) -> bool {
|
||||
filename.ends_with("tokenizer.json")
|
||||
|| filename.ends_with("tokenizer_config.json")
|
||||
|| filename.ends_with("special_tokens_map.json")
|
||||
|| filename.ends_with("vocab.json")
|
||||
|| filename.ends_with("merges.txt")
|
||||
|| filename.ends_with(".model") // SentencePiece models
|
||||
|| filename.ends_with(".tiktoken")
|
||||
|| is_chat_template_file(filename) // Include chat template files
|
||||
}
|
||||
|
||||
/// Checks if a file is a chat template file
|
||||
fn is_chat_template_file(filename: &str) -> bool {
|
||||
filename.ends_with(".jinja") // Direct Jinja files
|
||||
|| filename == "chat_template.json" // JSON file containing Jinja template
|
||||
}
|
||||
|
||||
/// Attempt to download tokenizer files from Hugging Face
|
||||
/// Returns the directory containing the downloaded tokenizer files
|
||||
pub async fn download_tokenizer_from_hf(model_id: impl AsRef<Path>) -> anyhow::Result<PathBuf> {
|
||||
let model_id = model_id.as_ref();
|
||||
let api = ApiBuilder::from_env().with_progress(true).build()?;
|
||||
let model_name = model_id.display().to_string();
|
||||
|
||||
let repo = api.model(model_name.clone());
|
||||
|
||||
let info = match repo.info().await {
|
||||
Ok(info) => info,
|
||||
Err(e) => {
|
||||
return Err(anyhow::anyhow!(
|
||||
"Failed to fetch model '{}' from HuggingFace: {}. Is this a valid HuggingFace ID?",
|
||||
model_name,
|
||||
e
|
||||
));
|
||||
}
|
||||
};
|
||||
|
||||
if info.siblings.is_empty() {
|
||||
return Err(anyhow::anyhow!(
|
||||
"Model '{}' exists but contains no downloadable files.",
|
||||
model_name
|
||||
));
|
||||
}
|
||||
|
||||
let mut cache_dir = None;
|
||||
let mut tokenizer_files_found = false;
|
||||
|
||||
// First, identify all tokenizer files to download
|
||||
let tokenizer_files: Vec<_> = info
|
||||
.siblings
|
||||
.iter()
|
||||
.filter(|sib| {
|
||||
!IGNORED.contains(&sib.rfilename.as_str())
|
||||
&& !is_image(&sib.rfilename)
|
||||
&& !is_weight_file(&sib.rfilename)
|
||||
&& is_tokenizer_file(&sib.rfilename)
|
||||
})
|
||||
.collect();
|
||||
|
||||
if tokenizer_files.is_empty() {
|
||||
return Err(anyhow::anyhow!(
|
||||
"No tokenizer files found for model '{}'.",
|
||||
model_name
|
||||
));
|
||||
}
|
||||
|
||||
// Download all tokenizer files
|
||||
for sib in tokenizer_files {
|
||||
match repo.get(&sib.rfilename).await {
|
||||
Ok(path) => {
|
||||
if cache_dir.is_none() {
|
||||
cache_dir = path.parent().map(|p| p.to_path_buf());
|
||||
}
|
||||
tokenizer_files_found = true;
|
||||
}
|
||||
Err(e) => {
|
||||
return Err(anyhow::anyhow!(
|
||||
"Failed to download tokenizer file '{}' from model '{}': {}",
|
||||
sib.rfilename,
|
||||
model_name,
|
||||
e
|
||||
));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if !tokenizer_files_found {
|
||||
return Err(anyhow::anyhow!(
|
||||
"No tokenizer files could be downloaded for model '{}'.",
|
||||
model_name
|
||||
));
|
||||
}
|
||||
|
||||
match cache_dir {
|
||||
Some(dir) => {
|
||||
// Ensure we return the correct model directory, not a subfolder
|
||||
// Some models have an "original" subfolder for PyTorch weights
|
||||
// We want the main model directory that contains tokenizer files
|
||||
let final_dir = resolve_model_cache_dir(&dir, &model_name);
|
||||
Ok(final_dir)
|
||||
}
|
||||
None => Err(anyhow::anyhow!(
|
||||
"Invalid HF cache path for model '{}'",
|
||||
model_name
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
/// Attempt to download a model from Hugging Face (including weights)
|
||||
/// Returns the directory it is in
|
||||
/// If ignore_weights is true, model weight files will be skipped
|
||||
pub async fn from_hf(name: impl AsRef<Path>, ignore_weights: bool) -> anyhow::Result<PathBuf> {
|
||||
let name = name.as_ref();
|
||||
let api = ApiBuilder::from_env().with_progress(true).build()?;
|
||||
let model_name = name.display().to_string();
|
||||
|
||||
let repo = api.model(model_name.clone());
|
||||
|
||||
let info = match repo.info().await {
|
||||
Ok(info) => info,
|
||||
Err(e) => {
|
||||
return Err(anyhow::anyhow!(
|
||||
"Failed to fetch model '{}' from HuggingFace: {}. Is this a valid HuggingFace ID?",
|
||||
model_name,
|
||||
e
|
||||
));
|
||||
}
|
||||
};
|
||||
|
||||
if info.siblings.is_empty() {
|
||||
return Err(anyhow::anyhow!(
|
||||
"Model '{}' exists but contains no downloadable files.",
|
||||
model_name
|
||||
));
|
||||
}
|
||||
|
||||
let mut p = PathBuf::new();
|
||||
let mut files_downloaded = false;
|
||||
|
||||
for sib in info.siblings {
|
||||
if IGNORED.contains(&sib.rfilename.as_str()) || is_image(&sib.rfilename) {
|
||||
continue;
|
||||
}
|
||||
|
||||
// If ignore_weights is true, skip weight files
|
||||
if ignore_weights && is_weight_file(&sib.rfilename) {
|
||||
continue;
|
||||
}
|
||||
|
||||
match repo.get(&sib.rfilename).await {
|
||||
Ok(path) => {
|
||||
p = path;
|
||||
files_downloaded = true;
|
||||
}
|
||||
Err(e) => {
|
||||
return Err(anyhow::anyhow!(
|
||||
"Failed to download file '{}' from model '{}': {}",
|
||||
sib.rfilename,
|
||||
model_name,
|
||||
e
|
||||
));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if !files_downloaded {
|
||||
let file_type = if ignore_weights {
|
||||
"non-weight"
|
||||
} else {
|
||||
"valid"
|
||||
};
|
||||
return Err(anyhow::anyhow!(
|
||||
"No {} files found for model '{}'.",
|
||||
file_type,
|
||||
model_name
|
||||
));
|
||||
}
|
||||
|
||||
match p.parent() {
|
||||
Some(p) => {
|
||||
let final_dir = resolve_model_cache_dir(p, &model_name);
|
||||
Ok(final_dir)
|
||||
}
|
||||
None => Err(anyhow::anyhow!("Invalid HF cache path: {}", p.display())),
|
||||
}
|
||||
}
|
||||
|
||||
/// Resolve the correct model cache directory
|
||||
/// Handles cases where files might be in subfolders (e.g., "original" folder)
|
||||
fn resolve_model_cache_dir(path: &Path, model_name: &str) -> PathBuf {
|
||||
// Check if we're in a subfolder like "original"
|
||||
if let Some(parent) = path.parent() {
|
||||
if let Some(folder_name) = path.file_name() {
|
||||
if folder_name == "original" {
|
||||
// We're in the "original" subfolder, go up one level
|
||||
return parent.to_path_buf();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Check if the current path contains the model name components
|
||||
// This helps ensure we're at the right directory level
|
||||
let model_parts: Vec<&str> = model_name.split('/').collect();
|
||||
if model_parts.len() >= 2 {
|
||||
let expected_pattern = format!(
|
||||
"models--{}--{}",
|
||||
model_parts[0].replace("-", "--"),
|
||||
model_parts[1].replace("-", "--")
|
||||
);
|
||||
|
||||
if path.to_string_lossy().contains(&expected_pattern) {
|
||||
// We're already at the correct level
|
||||
return path.to_path_buf();
|
||||
}
|
||||
|
||||
let mut current = path.to_path_buf();
|
||||
|
||||
// First check if current path already contains tokenizer files
|
||||
if current.join("tokenizer.json").exists() || current.join("tokenizer_config.json").exists()
|
||||
{
|
||||
return current;
|
||||
}
|
||||
|
||||
// If not, traverse up to find the model root, then look in snapshots
|
||||
while let Some(parent) = current.parent() {
|
||||
if parent.to_string_lossy().contains(&expected_pattern) {
|
||||
let snapshots_dir = parent.join("snapshots");
|
||||
if snapshots_dir.exists() && snapshots_dir.is_dir() {
|
||||
if let Ok(entries) = std::fs::read_dir(&snapshots_dir) {
|
||||
for entry in entries.flatten() {
|
||||
let snapshot_path = entry.path();
|
||||
if snapshot_path.is_dir()
|
||||
&& (snapshot_path.join("tokenizer.json").exists()
|
||||
|| snapshot_path.join("tokenizer_config.json").exists())
|
||||
{
|
||||
return snapshot_path;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
return parent.to_path_buf();
|
||||
}
|
||||
current = parent.to_path_buf();
|
||||
}
|
||||
}
|
||||
|
||||
path.to_path_buf()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_is_tokenizer_file() {
|
||||
assert!(is_tokenizer_file("tokenizer.json"));
|
||||
assert!(is_tokenizer_file("tokenizer_config.json"));
|
||||
assert!(is_tokenizer_file("special_tokens_map.json"));
|
||||
assert!(is_tokenizer_file("vocab.json"));
|
||||
assert!(is_tokenizer_file("merges.txt"));
|
||||
assert!(is_tokenizer_file("spiece.model"));
|
||||
assert!(is_tokenizer_file("chat_template.jinja"));
|
||||
assert!(is_tokenizer_file("template.jinja"));
|
||||
assert!(!is_tokenizer_file("model.bin"));
|
||||
assert!(!is_tokenizer_file("README.md"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_is_chat_template_file() {
|
||||
assert!(is_chat_template_file("chat_template.jinja"));
|
||||
assert!(is_chat_template_file("template.jinja"));
|
||||
assert!(is_chat_template_file("any_file.jinja"));
|
||||
assert!(is_chat_template_file("chat_template.json"));
|
||||
assert!(!is_chat_template_file("tokenizer.json"));
|
||||
assert!(!is_chat_template_file("other_file.json"));
|
||||
assert!(!is_chat_template_file("chat_template"));
|
||||
assert!(!is_chat_template_file("README.md"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_is_weight_file() {
|
||||
assert!(is_weight_file("model.bin"));
|
||||
assert!(is_weight_file("model.safetensors"));
|
||||
assert!(is_weight_file("pytorch_model.bin"));
|
||||
assert!(!is_weight_file("tokenizer.json"));
|
||||
assert!(!is_weight_file("config.json"));
|
||||
}
|
||||
}
|
||||
@@ -1,359 +0,0 @@
|
||||
use std::collections::HashMap;
|
||||
|
||||
use anyhow::{Error, Result};
|
||||
use tokenizers::{processors::template::TemplateProcessing, tokenizer::Tokenizer as HfTokenizer};
|
||||
use tracing::debug;
|
||||
|
||||
use super::{
|
||||
chat_template::{
|
||||
detect_chat_template_content_format, ChatTemplateContentFormat, ChatTemplateParams,
|
||||
ChatTemplateProcessor,
|
||||
},
|
||||
traits::{Decoder, Encoder, Encoding, SpecialTokens, TokenIdType, Tokenizer as TokenizerTrait},
|
||||
};
|
||||
|
||||
/// HuggingFace tokenizer wrapper
|
||||
pub struct HuggingFaceTokenizer {
|
||||
tokenizer: HfTokenizer,
|
||||
special_tokens: SpecialTokens,
|
||||
vocab: HashMap<String, TokenIdType>,
|
||||
reverse_vocab: HashMap<TokenIdType, String>,
|
||||
chat_template: Option<String>,
|
||||
/// Detected chat template content format (computed once at initialization)
|
||||
content_format: ChatTemplateContentFormat,
|
||||
}
|
||||
|
||||
impl HuggingFaceTokenizer {
|
||||
/// Create a tokenizer from a HuggingFace tokenizer JSON file
|
||||
pub fn from_file(file_path: &str) -> Result<Self> {
|
||||
// Try to auto-discover chat template if not explicitly provided
|
||||
let path = std::path::Path::new(file_path);
|
||||
let chat_template_path = path
|
||||
.parent()
|
||||
.and_then(crate::tokenizer::factory::discover_chat_template_in_dir);
|
||||
Self::from_file_with_chat_template(file_path, chat_template_path.as_deref())
|
||||
}
|
||||
|
||||
/// Create a tokenizer from a HuggingFace tokenizer JSON file with an optional chat template
|
||||
pub fn from_file_with_chat_template(
|
||||
file_path: &str,
|
||||
chat_template_path: Option<&str>,
|
||||
) -> Result<Self> {
|
||||
let mut tokenizer = HfTokenizer::from_file(file_path)
|
||||
.map_err(|e| Error::msg(format!("Failed to load tokenizer: {}", e)))?;
|
||||
|
||||
// Extract special tokens
|
||||
let special_tokens = Self::extract_special_tokens(&tokenizer);
|
||||
|
||||
// Build vocab mappings (include special tokens to get added_tokens like <|im_start|>)
|
||||
let vocab = tokenizer.get_vocab(true); // true = include special tokens and added_tokens
|
||||
let reverse_vocab: HashMap<TokenIdType, String> = vocab
|
||||
.iter()
|
||||
.map(|(token, &id)| (id, token.clone()))
|
||||
.collect();
|
||||
|
||||
// Load chat template and tokenizer config
|
||||
let (chat_template, add_bos_token, add_eos_token) =
|
||||
if let Some(template_path) = chat_template_path {
|
||||
// Load from specified .jinja file
|
||||
(
|
||||
Self::load_chat_template_from_file(template_path)?,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
} else {
|
||||
// Try to load from tokenizer_config.json
|
||||
Self::load_chat_template_and_config(file_path)
|
||||
};
|
||||
|
||||
// Detect content format once at initialization
|
||||
let content_format = if let Some(ref template) = chat_template {
|
||||
detect_chat_template_content_format(template)
|
||||
} else {
|
||||
ChatTemplateContentFormat::String // Default if no template
|
||||
};
|
||||
|
||||
// Configure post_processor based on tokenizer_config.json (matches Python transformers)
|
||||
// Only modify when at least one setting is explicitly true
|
||||
let needs_eos = add_eos_token == Some(true);
|
||||
let needs_bos = match add_bos_token {
|
||||
Some(true) => true,
|
||||
Some(false) => false,
|
||||
// Not set: preserve existing behavior from tokenizer.json
|
||||
None => needs_eos && Self::tokenizer_adds_special_tokens(&tokenizer),
|
||||
};
|
||||
|
||||
if needs_bos || needs_eos {
|
||||
if let Some(post_processor) =
|
||||
Self::build_post_processor(needs_bos, needs_eos, &special_tokens, &vocab)
|
||||
{
|
||||
debug!(needs_bos, needs_eos, "Configured post_processor");
|
||||
tokenizer.with_post_processor(Some(post_processor));
|
||||
}
|
||||
}
|
||||
|
||||
Ok(HuggingFaceTokenizer {
|
||||
tokenizer,
|
||||
special_tokens,
|
||||
vocab,
|
||||
reverse_vocab,
|
||||
chat_template,
|
||||
content_format,
|
||||
})
|
||||
}
|
||||
|
||||
/// Check if the tokenizer's post_processor adds special tokens (e.g., BOS)
|
||||
fn tokenizer_adds_special_tokens(tokenizer: &HfTokenizer) -> bool {
|
||||
tokenizer
|
||||
.encode("", true)
|
||||
.map(|enc| !enc.get_ids().is_empty())
|
||||
.unwrap_or(false)
|
||||
}
|
||||
|
||||
/// Build a TemplateProcessing post_processor (matches Python transformers' update_post_processor)
|
||||
/// Template format: "{bos}:0 $A:0 {eos}:0" with optional BOS/EOS based on config
|
||||
fn build_post_processor(
|
||||
add_bos_token: bool,
|
||||
add_eos_token: bool,
|
||||
special_tokens: &SpecialTokens,
|
||||
vocab: &HashMap<String, TokenIdType>,
|
||||
) -> Option<TemplateProcessing> {
|
||||
// Build template string exactly like Python:
|
||||
// single = f"{(bos + ':0 ') if add_bos_token else ''}$A:0{(' ' + eos + ':0') if add_eos_token else ''}"
|
||||
let mut template = String::with_capacity(32);
|
||||
let mut tokens = Vec::with_capacity(2);
|
||||
|
||||
if add_bos_token {
|
||||
let bos = special_tokens.bos_token.as_ref()?;
|
||||
let bos_id = vocab.get(bos).copied()?;
|
||||
template.push_str(bos);
|
||||
template.push_str(":0 ");
|
||||
tokens.push((bos.clone(), bos_id));
|
||||
}
|
||||
|
||||
template.push_str("$A:0");
|
||||
|
||||
if add_eos_token {
|
||||
let eos = special_tokens.eos_token.as_ref()?;
|
||||
let eos_id = vocab.get(eos).copied()?;
|
||||
template.push(' ');
|
||||
template.push_str(eos);
|
||||
template.push_str(":0");
|
||||
tokens.push((eos.clone(), eos_id));
|
||||
}
|
||||
|
||||
TemplateProcessing::builder()
|
||||
.try_single(template.as_str())
|
||||
.ok()?
|
||||
.special_tokens(tokens)
|
||||
.build()
|
||||
.ok()
|
||||
}
|
||||
|
||||
/// Create from an existing HuggingFace tokenizer
|
||||
pub fn from_tokenizer(tokenizer: HfTokenizer) -> Self {
|
||||
let special_tokens = Self::extract_special_tokens(&tokenizer);
|
||||
let vocab = tokenizer.get_vocab(true); // true = include special tokens and added_tokens
|
||||
let reverse_vocab: HashMap<TokenIdType, String> = vocab
|
||||
.iter()
|
||||
.map(|(token, &id)| (id, token.clone()))
|
||||
.collect();
|
||||
|
||||
HuggingFaceTokenizer {
|
||||
tokenizer,
|
||||
special_tokens,
|
||||
vocab,
|
||||
reverse_vocab,
|
||||
chat_template: None,
|
||||
content_format: ChatTemplateContentFormat::String, // Default
|
||||
}
|
||||
}
|
||||
|
||||
/// Extract special tokens from the tokenizer
|
||||
fn extract_special_tokens(tokenizer: &HfTokenizer) -> SpecialTokens {
|
||||
// Get vocab with special tokens included (added_tokens like <|im_start|>)
|
||||
let vocab = tokenizer.get_vocab(true);
|
||||
|
||||
let find_token = |patterns: &[&str]| -> Option<String> {
|
||||
for pattern in patterns {
|
||||
if vocab.contains_key(*pattern) {
|
||||
return Some(pattern.to_string());
|
||||
}
|
||||
}
|
||||
None
|
||||
};
|
||||
|
||||
// Extract additional special tokens using the tokenizers library API
|
||||
let additional_special_tokens: Vec<String> = tokenizer
|
||||
.get_added_tokens_decoder()
|
||||
.iter()
|
||||
.filter(|(_id, token)| token.special) // Only tokens marked as special: true
|
||||
.map(|(_id, token)| token.content.clone())
|
||||
.collect();
|
||||
|
||||
SpecialTokens {
|
||||
bos_token: find_token(&["<s>", "<|startoftext|>", "<BOS>", "[CLS]"]),
|
||||
eos_token: find_token(&["</s>", "<|endoftext|>", "<EOS>", "[SEP]"]),
|
||||
unk_token: find_token(&["<unk>", "<UNK>", "[UNK]"]),
|
||||
sep_token: find_token(&["[SEP]", "<sep>", "<SEP>"]),
|
||||
pad_token: find_token(&["<pad>", "<PAD>", "[PAD]"]),
|
||||
cls_token: find_token(&["[CLS]", "<cls>", "<CLS>"]),
|
||||
mask_token: find_token(&["[MASK]", "<mask>", "<MASK>"]),
|
||||
additional_special_tokens,
|
||||
}
|
||||
}
|
||||
|
||||
/// Load chat template and special token settings from tokenizer_config.json
|
||||
/// Returns Option<bool> to distinguish between explicit false vs not set
|
||||
fn load_chat_template_and_config(
|
||||
tokenizer_path: &str,
|
||||
) -> (Option<String>, Option<bool>, Option<bool>) {
|
||||
(|| {
|
||||
let path = std::path::Path::new(tokenizer_path);
|
||||
let config_path = path.parent()?.join("tokenizer_config.json");
|
||||
|
||||
if !config_path.exists() {
|
||||
return None;
|
||||
}
|
||||
|
||||
let config_str = config_path.to_str()?;
|
||||
let content = std::fs::read_to_string(&config_path).ok()?;
|
||||
let config: serde_json::Value = serde_json::from_str(&content).ok()?;
|
||||
|
||||
let chat_template = super::chat_template::load_chat_template_from_config(config_str)
|
||||
.ok()
|
||||
.flatten();
|
||||
|
||||
let add_bos_token = config.get("add_bos_token").and_then(|v| v.as_bool());
|
||||
let add_eos_token = config.get("add_eos_token").and_then(|v| v.as_bool());
|
||||
|
||||
Some((chat_template, add_bos_token, add_eos_token))
|
||||
})()
|
||||
.unwrap_or((None, None, None))
|
||||
}
|
||||
|
||||
/// Load chat template from a file (.jinja or .json containing Jinja)
|
||||
fn load_chat_template_from_file(template_path: &str) -> Result<Option<String>> {
|
||||
use std::fs;
|
||||
|
||||
let content = fs::read_to_string(template_path)
|
||||
.map_err(|e| Error::msg(format!("Failed to read chat template file: {}", e)))?;
|
||||
|
||||
// Check if it's a JSON file containing a Jinja template
|
||||
if template_path.ends_with(".json") {
|
||||
// Parse JSON and extract the template string
|
||||
let json_value: serde_json::Value = serde_json::from_str(&content)
|
||||
.map_err(|e| Error::msg(format!("Failed to parse chat_template.json: {}", e)))?;
|
||||
|
||||
if let Some(template_str) = json_value.as_str() {
|
||||
return Ok(Some(template_str.to_string()));
|
||||
} else if let Some(obj) = json_value.as_object() {
|
||||
if let Some(template_value) = obj.get("chat_template") {
|
||||
if let Some(template_str) = template_value.as_str() {
|
||||
return Ok(Some(template_str.to_string()));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return Err(Error::msg(
|
||||
"chat_template.json does not contain a valid template",
|
||||
));
|
||||
}
|
||||
|
||||
// Otherwise it's a plain .jinja file
|
||||
// Clean up the template (similar to Python implementation)
|
||||
let template = content.trim().replace("\\n", "\n");
|
||||
|
||||
Ok(Some(template))
|
||||
}
|
||||
|
||||
/// Set or override the chat template
|
||||
pub fn set_chat_template(&mut self, template: String) {
|
||||
// Detect format for the new template
|
||||
self.content_format = detect_chat_template_content_format(&template);
|
||||
self.chat_template = Some(template);
|
||||
}
|
||||
|
||||
/// Get the content format expected by the chat template
|
||||
pub fn chat_template_content_format(&self) -> ChatTemplateContentFormat {
|
||||
self.content_format
|
||||
}
|
||||
|
||||
/// Apply chat template if available
|
||||
///
|
||||
/// Takes transformed JSON Values (already transformed based on content format)
|
||||
pub fn apply_chat_template(
|
||||
&self,
|
||||
messages: &[serde_json::Value],
|
||||
params: ChatTemplateParams,
|
||||
) -> Result<String> {
|
||||
if let Some(ref template) = self.chat_template {
|
||||
let processor = ChatTemplateProcessor::new(template.clone());
|
||||
processor.apply_chat_template(messages, params)
|
||||
} else {
|
||||
Err(Error::msg(
|
||||
"Cannot use chat template functions because tokenizer.chat_template is not set and no template \
|
||||
argument was passed! For information about writing templates and setting the \
|
||||
tokenizer.chat_template attribute, please see the documentation at \
|
||||
https://huggingface.co/docs/transformers/main/en/chat_templating"
|
||||
))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Encoder for HuggingFaceTokenizer {
|
||||
fn encode(&self, input: &str, add_special_tokens: bool) -> Result<Encoding> {
|
||||
self.tokenizer
|
||||
.encode(input, add_special_tokens)
|
||||
.map_err(|e| Error::msg(format!("Encoding failed: {}", e)))
|
||||
.map(|encoding| Encoding::Hf(Box::new(encoding)))
|
||||
}
|
||||
|
||||
fn encode_batch(&self, inputs: &[&str], add_special_tokens: bool) -> Result<Vec<Encoding>> {
|
||||
self.tokenizer
|
||||
.encode_batch(inputs.to_vec(), add_special_tokens)
|
||||
.map_err(|e| Error::msg(format!("Batch encoding failed: {}", e)))
|
||||
.map(|encodings| {
|
||||
encodings
|
||||
.into_iter()
|
||||
.map(|e| Encoding::Hf(Box::new(e)))
|
||||
.collect()
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl Decoder for HuggingFaceTokenizer {
|
||||
fn decode(&self, token_ids: &[TokenIdType], skip_special_tokens: bool) -> Result<String> {
|
||||
self.tokenizer
|
||||
.decode(token_ids, skip_special_tokens)
|
||||
.map_err(|e| Error::msg(format!("Decoding failed: {}", e)))
|
||||
}
|
||||
}
|
||||
|
||||
impl TokenizerTrait for HuggingFaceTokenizer {
|
||||
fn vocab_size(&self) -> usize {
|
||||
self.tokenizer.get_vocab_size(false)
|
||||
}
|
||||
|
||||
fn get_special_tokens(&self) -> &SpecialTokens {
|
||||
&self.special_tokens
|
||||
}
|
||||
|
||||
fn token_to_id(&self, token: &str) -> Option<TokenIdType> {
|
||||
self.vocab.get(token).copied()
|
||||
}
|
||||
|
||||
fn id_to_token(&self, id: TokenIdType) -> Option<String> {
|
||||
self.reverse_vocab.get(&id).cloned()
|
||||
}
|
||||
|
||||
fn as_any(&self) -> &dyn std::any::Any {
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
// Note: Actual tokenizer tests would require a real tokenizer file
|
||||
// These would be integration tests rather than unit tests
|
||||
}
|
||||
@@ -1,128 +0,0 @@
|
||||
//! Mock tokenizer implementation for testing
|
||||
|
||||
use std::collections::HashMap;
|
||||
|
||||
use anyhow::Result;
|
||||
|
||||
use super::traits::{Decoder, Encoder, Encoding, SpecialTokens, Tokenizer as TokenizerTrait};
|
||||
|
||||
/// Mock tokenizer for testing purposes
|
||||
pub struct MockTokenizer {
|
||||
vocab: HashMap<String, u32>,
|
||||
reverse_vocab: HashMap<u32, String>,
|
||||
special_tokens: SpecialTokens,
|
||||
}
|
||||
|
||||
impl Default for MockTokenizer {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
impl MockTokenizer {
|
||||
pub fn new() -> Self {
|
||||
let mut vocab = HashMap::new();
|
||||
let mut reverse_vocab = HashMap::new();
|
||||
|
||||
// Add some basic tokens
|
||||
let tokens = vec![
|
||||
("Hello", 1),
|
||||
("world", 2),
|
||||
("test", 3),
|
||||
("token", 4),
|
||||
(" ", 5),
|
||||
(".", 6),
|
||||
("<eos>", 999),
|
||||
("<bos>", 1000),
|
||||
("<|im_start|>", 1001),
|
||||
("<|im_end|>", 1002),
|
||||
("<|eot_id|>", 1003),
|
||||
("system", 7),
|
||||
("user", 8),
|
||||
("assistant", 9),
|
||||
];
|
||||
|
||||
for (token, id) in tokens {
|
||||
vocab.insert(token.to_string(), id);
|
||||
reverse_vocab.insert(id, token.to_string());
|
||||
}
|
||||
|
||||
let special_tokens = SpecialTokens {
|
||||
bos_token: Some("<bos>".to_string()),
|
||||
eos_token: Some("<eos>".to_string()),
|
||||
unk_token: Some("<unk>".to_string()),
|
||||
sep_token: None,
|
||||
pad_token: None,
|
||||
cls_token: None,
|
||||
mask_token: None,
|
||||
additional_special_tokens: vec![],
|
||||
};
|
||||
|
||||
Self {
|
||||
vocab,
|
||||
reverse_vocab,
|
||||
special_tokens,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Encoder for MockTokenizer {
|
||||
fn encode(&self, input: &str, _add_special_tokens: bool) -> Result<Encoding> {
|
||||
// Simple word-based tokenization using the vocab
|
||||
// Split by whitespace and look up each word (decoder adds spaces back)
|
||||
let tokens: Vec<u32> = input
|
||||
.split_whitespace()
|
||||
.filter_map(|word| self.vocab.get(word).copied())
|
||||
.collect();
|
||||
|
||||
Ok(Encoding::Sp(tokens))
|
||||
}
|
||||
|
||||
fn encode_batch(&self, inputs: &[&str], add_special_tokens: bool) -> Result<Vec<Encoding>> {
|
||||
inputs
|
||||
.iter()
|
||||
.map(|input| self.encode(input, add_special_tokens))
|
||||
.collect()
|
||||
}
|
||||
}
|
||||
|
||||
impl Decoder for MockTokenizer {
|
||||
fn decode(&self, token_ids: &[u32], skip_special_tokens: bool) -> Result<String> {
|
||||
let tokens: Vec<String> = token_ids
|
||||
.iter()
|
||||
.filter_map(|id| {
|
||||
self.reverse_vocab.get(id).and_then(|token| {
|
||||
if skip_special_tokens && (token == "<eos>" || token == "<bos>") {
|
||||
None
|
||||
} else {
|
||||
Some(token.clone())
|
||||
}
|
||||
})
|
||||
})
|
||||
.collect();
|
||||
|
||||
Ok(tokens.join(" "))
|
||||
}
|
||||
}
|
||||
|
||||
impl TokenizerTrait for MockTokenizer {
|
||||
fn vocab_size(&self) -> usize {
|
||||
self.vocab.len()
|
||||
}
|
||||
|
||||
fn get_special_tokens(&self) -> &SpecialTokens {
|
||||
&self.special_tokens
|
||||
}
|
||||
|
||||
fn token_to_id(&self, token: &str) -> Option<u32> {
|
||||
self.vocab.get(token).copied()
|
||||
}
|
||||
|
||||
fn id_to_token(&self, id: u32) -> Option<String> {
|
||||
self.reverse_vocab.get(&id).cloned()
|
||||
}
|
||||
|
||||
fn as_any(&self) -> &dyn std::any::Any {
|
||||
self
|
||||
}
|
||||
}
|
||||
@@ -1,124 +0,0 @@
|
||||
use std::{ops::Deref, sync::Arc};
|
||||
|
||||
use anyhow::Result;
|
||||
|
||||
pub mod cache;
|
||||
pub mod factory;
|
||||
pub mod hub;
|
||||
pub mod mock;
|
||||
pub mod registry;
|
||||
pub mod sequence;
|
||||
pub mod stop;
|
||||
pub mod stream;
|
||||
pub mod traits;
|
||||
|
||||
// Feature-gated modules
|
||||
|
||||
pub mod chat_template;
|
||||
|
||||
pub mod huggingface;
|
||||
|
||||
pub mod tiktoken;
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests;
|
||||
|
||||
// Internal imports for Tokenizer struct
|
||||
use factory::{create_tokenizer_from_file, create_tokenizer_with_chat_template};
|
||||
// Re-export types used outside this module
|
||||
pub use huggingface::HuggingFaceTokenizer;
|
||||
pub use registry::{LoadError, LoadOutcome, TokenizerRegistry};
|
||||
pub use stop::StopSequenceDecoder;
|
||||
pub use stream::DecodeStream;
|
||||
pub use traits::{Decoder, Encoder, Encoding, SpecialTokens};
|
||||
|
||||
/// Main tokenizer wrapper that provides a unified interface for different tokenizer implementations
|
||||
#[derive(Clone)]
|
||||
pub struct Tokenizer(Arc<dyn traits::Tokenizer>);
|
||||
|
||||
impl Tokenizer {
|
||||
/// Create a tokenizer from a file path
|
||||
pub fn from_file(file_path: &str) -> Result<Tokenizer> {
|
||||
Ok(Tokenizer(create_tokenizer_from_file(file_path)?))
|
||||
}
|
||||
|
||||
/// Create a tokenizer from a file path with an optional chat template
|
||||
pub fn from_file_with_chat_template(
|
||||
file_path: &str,
|
||||
chat_template_path: Option<&str>,
|
||||
) -> Result<Tokenizer> {
|
||||
Ok(Tokenizer(create_tokenizer_with_chat_template(
|
||||
file_path,
|
||||
chat_template_path,
|
||||
)?))
|
||||
}
|
||||
|
||||
/// Create a tokenizer from an Arc<dyn Tokenizer>
|
||||
pub fn from_arc(tokenizer: Arc<dyn traits::Tokenizer>) -> Self {
|
||||
Tokenizer(tokenizer)
|
||||
}
|
||||
|
||||
/// Create a stateful sequence object for decoding token_ids into text
|
||||
pub fn decode_stream(
|
||||
&self,
|
||||
prompt_token_ids: &[u32],
|
||||
skip_special_tokens: bool,
|
||||
) -> DecodeStream {
|
||||
DecodeStream::new(self.0.clone(), prompt_token_ids, skip_special_tokens)
|
||||
}
|
||||
|
||||
/// Direct encode method
|
||||
///
|
||||
/// Set `add_special_tokens` to `true` for embeddings (to add BOS/EOS tokens configured in tokenizer_config.json),
|
||||
/// or `false` for chat completion (where the chat template handles special tokens).
|
||||
pub fn encode(&self, input: &str, add_special_tokens: bool) -> Result<Encoding> {
|
||||
self.0.encode(input, add_special_tokens)
|
||||
}
|
||||
|
||||
/// Direct batch encode method
|
||||
///
|
||||
/// Set `add_special_tokens` to `true` for embeddings (to add BOS/EOS tokens configured in tokenizer_config.json),
|
||||
/// or `false` for chat completion (where the chat template handles special tokens).
|
||||
pub fn encode_batch(&self, inputs: &[&str], add_special_tokens: bool) -> Result<Vec<Encoding>> {
|
||||
self.0.encode_batch(inputs, add_special_tokens)
|
||||
}
|
||||
|
||||
/// Direct decode method
|
||||
pub fn decode(&self, token_ids: &[u32], skip_special_tokens: bool) -> Result<String> {
|
||||
self.0.decode(token_ids, skip_special_tokens)
|
||||
}
|
||||
|
||||
/// Get vocabulary size
|
||||
pub fn vocab_size(&self) -> usize {
|
||||
self.0.vocab_size()
|
||||
}
|
||||
|
||||
/// Get special tokens
|
||||
pub fn get_special_tokens(&self) -> &SpecialTokens {
|
||||
self.0.get_special_tokens()
|
||||
}
|
||||
|
||||
/// Convert token string to ID
|
||||
pub fn token_to_id(&self, token: &str) -> Option<u32> {
|
||||
self.0.token_to_id(token)
|
||||
}
|
||||
|
||||
/// Convert ID to token string
|
||||
pub fn id_to_token(&self, id: u32) -> Option<String> {
|
||||
self.0.id_to_token(id)
|
||||
}
|
||||
}
|
||||
|
||||
impl Deref for Tokenizer {
|
||||
type Target = Arc<dyn traits::Tokenizer>;
|
||||
|
||||
fn deref(&self) -> &Self::Target {
|
||||
&self.0
|
||||
}
|
||||
}
|
||||
|
||||
impl From<Arc<dyn traits::Tokenizer>> for Tokenizer {
|
||||
fn from(tokenizer: Arc<dyn traits::Tokenizer>) -> Self {
|
||||
Tokenizer(tokenizer)
|
||||
}
|
||||
}
|
||||
@@ -1,677 +0,0 @@
|
||||
//! Tokenizer Registry for dynamic tokenizer loading
|
||||
//!
|
||||
//! Provides thread-safe, deduplicated tokenizer loading for IGW mode where
|
||||
//! multiple routers (HTTP and gRPC) need to share tokenizers across workers.
|
||||
//!
|
||||
//! ## ID vs Name Lookup
|
||||
//!
|
||||
//! Tokenizers are stored with two keys:
|
||||
//! - **ID (UUID)**: Unique identifier generated at registration, immutable
|
||||
//! - **Name**: User-provided identifier, must be unique
|
||||
//!
|
||||
//! Lookup behavior:
|
||||
//! - `get(key)`: Tries name first, then ID (backward compatible)
|
||||
//! - `get_by_id(id)`: Exact ID match only
|
||||
//! - `get_by_name(name)`: Exact name match only
|
||||
//! - `remove(name)`: Removes by name
|
||||
//! - `remove_by_id(id)`: Removes by ID
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use dashmap::DashMap;
|
||||
use thiserror::Error;
|
||||
use tokio::sync::Mutex;
|
||||
use tracing::{debug, info};
|
||||
use uuid::Uuid;
|
||||
|
||||
use super::traits::Tokenizer;
|
||||
|
||||
/// Outcome of a tokenizer load operation
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum LoadOutcome {
|
||||
/// Tokenizer was newly loaded and registered
|
||||
Loaded { id: String },
|
||||
/// Tokenizer already existed, returning existing ID
|
||||
AlreadyExists { id: String },
|
||||
}
|
||||
|
||||
impl LoadOutcome {
|
||||
/// Get the ID regardless of outcome
|
||||
pub fn id(&self) -> &str {
|
||||
match self {
|
||||
LoadOutcome::Loaded { id } => id,
|
||||
LoadOutcome::AlreadyExists { id } => id,
|
||||
}
|
||||
}
|
||||
|
||||
/// Returns true if the tokenizer was newly loaded
|
||||
pub fn is_newly_loaded(&self) -> bool {
|
||||
matches!(self, LoadOutcome::Loaded { .. })
|
||||
}
|
||||
}
|
||||
|
||||
/// Error type for tokenizer loading operations
|
||||
#[derive(Debug, Error)]
|
||||
pub enum LoadError {
|
||||
/// Name cannot be empty
|
||||
#[error("tokenizer name cannot be empty")]
|
||||
EmptyName,
|
||||
/// Source cannot be empty
|
||||
#[error("tokenizer source cannot be empty")]
|
||||
EmptySource,
|
||||
/// Loading failed
|
||||
#[error("{0}")]
|
||||
LoadFailed(String),
|
||||
}
|
||||
|
||||
/// Metadata and tokenizer instance for a registered tokenizer
|
||||
#[derive(Clone)]
|
||||
pub struct TokenizerEntry {
|
||||
/// Unique identifier (UUID)
|
||||
pub id: String,
|
||||
/// User-provided name
|
||||
pub name: String,
|
||||
/// Source path or HuggingFace model ID
|
||||
pub source: String,
|
||||
/// The tokenizer instance
|
||||
pub tokenizer: Arc<dyn Tokenizer>,
|
||||
}
|
||||
|
||||
/// Registry for managing tokenizers keyed by UUID
|
||||
///
|
||||
/// Features:
|
||||
/// - Thread-safe concurrent access using DashMap
|
||||
/// - Per-key locking to prevent duplicate loading
|
||||
/// - Lookup by UUID (primary) or name (secondary index)
|
||||
pub struct TokenizerRegistry {
|
||||
/// Storage for loaded tokenizers, keyed by UUID
|
||||
tokenizers: DashMap<String, TokenizerEntry>,
|
||||
/// Secondary index: name -> UUID for lookup
|
||||
name_to_id: DashMap<String, String>,
|
||||
/// Per-key locks to prevent duplicate loading
|
||||
loading_locks: DashMap<String, Arc<Mutex<()>>>,
|
||||
}
|
||||
|
||||
/// RAII guard that removes the loading lock entry on drop.
|
||||
/// Ensures cleanup happens on normal completion, early return, or panic.
|
||||
struct LoadingLockGuard<'a> {
|
||||
locks: &'a DashMap<String, Arc<Mutex<()>>>,
|
||||
key: String,
|
||||
}
|
||||
|
||||
impl Drop for LoadingLockGuard<'_> {
|
||||
fn drop(&mut self) {
|
||||
self.locks.remove(&self.key);
|
||||
}
|
||||
}
|
||||
|
||||
impl TokenizerRegistry {
|
||||
/// Create a new empty registry
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
tokenizers: DashMap::new(),
|
||||
name_to_id: DashMap::new(),
|
||||
loading_locks: DashMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Generate a new UUID for a tokenizer
|
||||
pub fn generate_id() -> String {
|
||||
Uuid::new_v4().to_string()
|
||||
}
|
||||
|
||||
/// Load and register a tokenizer
|
||||
///
|
||||
/// Validates inputs, handles deduplication, and loads the tokenizer if needed.
|
||||
/// Per-key locking ensures only one load happens per name, preventing race conditions.
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `id` - Pre-generated UUID (use `generate_id()` to create one)
|
||||
/// * `name` - User-provided name (used for deduplication, must not be empty)
|
||||
/// * `source` - Source path or HuggingFace model ID (must not be empty)
|
||||
/// * `loader` - Async function that loads the tokenizer
|
||||
///
|
||||
/// # Returns
|
||||
/// * `Ok(LoadOutcome::Loaded { id })` - Tokenizer was newly loaded
|
||||
/// * `Ok(LoadOutcome::AlreadyExists { id })` - Tokenizer already existed
|
||||
/// * `Err(LoadError)` - Validation failed or loading failed
|
||||
pub async fn load<F, Fut>(
|
||||
&self,
|
||||
id: &str,
|
||||
name: &str,
|
||||
source: &str,
|
||||
loader: F,
|
||||
) -> Result<LoadOutcome, LoadError>
|
||||
where
|
||||
F: FnOnce() -> Fut,
|
||||
Fut: std::future::Future<Output = Result<Arc<dyn Tokenizer>, String>>,
|
||||
{
|
||||
// Validate inputs
|
||||
if name.is_empty() {
|
||||
return Err(LoadError::EmptyName);
|
||||
}
|
||||
if source.is_empty() {
|
||||
return Err(LoadError::EmptySource);
|
||||
}
|
||||
|
||||
// Fast path: already loaded by name
|
||||
if let Some(existing_id) = self.name_to_id.get(name) {
|
||||
debug!("Tokenizer already registered for name: {}", name);
|
||||
return Ok(LoadOutcome::AlreadyExists {
|
||||
id: existing_id.clone(),
|
||||
});
|
||||
}
|
||||
|
||||
debug!("Tokenizer cache miss for name: {}", name);
|
||||
|
||||
// Acquire per-name lock to prevent duplicate loading
|
||||
let lock = self
|
||||
.loading_locks
|
||||
.entry(name.to_string())
|
||||
.or_insert_with(|| Arc::new(Mutex::new(())))
|
||||
.clone();
|
||||
|
||||
let _mutex_guard = lock.lock().await;
|
||||
let _lock_cleanup = LoadingLockGuard {
|
||||
locks: &self.loading_locks,
|
||||
key: name.to_string(),
|
||||
};
|
||||
|
||||
// Double-check after acquiring lock (another thread may have loaded it)
|
||||
if let Some(existing_id) = self.name_to_id.get(name) {
|
||||
debug!("Tokenizer loaded by another thread for name: {}", name);
|
||||
return Ok(LoadOutcome::AlreadyExists {
|
||||
id: existing_id.clone(),
|
||||
});
|
||||
}
|
||||
|
||||
// Load tokenizer
|
||||
info!("Loading tokenizer '{}' from source: {}", name, source);
|
||||
let result = loader().await;
|
||||
|
||||
let tokenizer = result.map_err(LoadError::LoadFailed)?;
|
||||
|
||||
// Create entry with provided ID
|
||||
let entry = TokenizerEntry {
|
||||
id: id.to_string(),
|
||||
name: name.to_string(),
|
||||
source: source.to_string(),
|
||||
tokenizer,
|
||||
};
|
||||
|
||||
// Store in registry
|
||||
self.tokenizers.insert(id.to_string(), entry);
|
||||
self.name_to_id.insert(name.to_string(), id.to_string());
|
||||
|
||||
info!(
|
||||
"Successfully registered tokenizer '{}' with id: {}",
|
||||
name, id
|
||||
);
|
||||
|
||||
Ok(LoadOutcome::Loaded { id: id.to_string() })
|
||||
}
|
||||
|
||||
/// Register a preloaded tokenizer with a pre-generated ID
|
||||
///
|
||||
/// Atomically inserts a tokenizer into the registry only if no tokenizer
|
||||
/// with the same name exists. Returns the ID if successful.
|
||||
///
|
||||
/// This is primarily used for testing. Production code should use `load()`.
|
||||
///
|
||||
/// # Returns
|
||||
/// * `Some(id)` - If the tokenizer was successfully registered
|
||||
/// * `None` - If a tokenizer with this name already existed
|
||||
#[cfg(test)]
|
||||
pub(crate) fn register(
|
||||
&self,
|
||||
id: &str,
|
||||
name: &str,
|
||||
source: &str,
|
||||
tokenizer: Arc<dyn Tokenizer>,
|
||||
) -> Option<String> {
|
||||
use dashmap::mapref::entry::Entry;
|
||||
|
||||
// Check if name already exists
|
||||
match self.name_to_id.entry(name.to_string()) {
|
||||
Entry::Occupied(_) => {
|
||||
debug!(
|
||||
"Tokenizer already exists for name: {}, skipping registration",
|
||||
name
|
||||
);
|
||||
None
|
||||
}
|
||||
Entry::Vacant(name_entry) => {
|
||||
let entry = TokenizerEntry {
|
||||
id: id.to_string(),
|
||||
name: name.to_string(),
|
||||
source: source.to_string(),
|
||||
tokenizer,
|
||||
};
|
||||
|
||||
info!("Registering tokenizer '{}' with id: {}", name, id);
|
||||
self.tokenizers.insert(id.to_string(), entry);
|
||||
name_entry.insert(id.to_string());
|
||||
Some(id.to_string())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Get a tokenizer by UUID
|
||||
pub fn get_by_id(&self, id: &str) -> Option<TokenizerEntry> {
|
||||
self.tokenizers.get(id).map(|e| e.clone())
|
||||
}
|
||||
|
||||
/// Get a tokenizer by name
|
||||
pub fn get_by_name(&self, name: &str) -> Option<TokenizerEntry> {
|
||||
self.name_to_id
|
||||
.get(name)
|
||||
.and_then(|id| self.tokenizers.get(id.as_str()).map(|e| e.clone()))
|
||||
}
|
||||
|
||||
/// Get a tokenizer (for backward compatibility, tries name first then ID)
|
||||
pub fn get(&self, name_or_id: &str) -> Option<Arc<dyn Tokenizer>> {
|
||||
self.get_by_name(name_or_id)
|
||||
.or_else(|| self.get_by_id(name_or_id))
|
||||
.map(|e| e.tokenizer)
|
||||
}
|
||||
|
||||
/// Check if a tokenizer is registered by name
|
||||
pub fn contains(&self, name: &str) -> bool {
|
||||
self.name_to_id.contains_key(name)
|
||||
}
|
||||
|
||||
/// Check if a tokenizer is registered by ID
|
||||
pub fn contains_id(&self, id: &str) -> bool {
|
||||
self.tokenizers.contains_key(id)
|
||||
}
|
||||
|
||||
/// Get the number of loaded tokenizers
|
||||
pub fn len(&self) -> usize {
|
||||
self.tokenizers.len()
|
||||
}
|
||||
|
||||
/// Check if the registry is empty
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.tokenizers.is_empty()
|
||||
}
|
||||
|
||||
/// List all registered tokenizers
|
||||
pub fn list(&self) -> Vec<TokenizerEntry> {
|
||||
let mut entries: Vec<TokenizerEntry> =
|
||||
self.tokenizers.iter().map(|e| e.value().clone()).collect();
|
||||
entries.sort_by(|a, b| a.name.cmp(&b.name));
|
||||
entries
|
||||
}
|
||||
|
||||
/// Remove a tokenizer by ID
|
||||
///
|
||||
/// Returns the entry if it was present.
|
||||
pub fn remove_by_id(&self, id: &str) -> Option<TokenizerEntry> {
|
||||
if let Some((_, entry)) = self.tokenizers.remove(id) {
|
||||
self.name_to_id.remove(&entry.name);
|
||||
Some(entry)
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
/// Remove a tokenizer by name
|
||||
///
|
||||
/// Returns the entry if it was present.
|
||||
pub fn remove(&self, name: &str) -> Option<TokenizerEntry> {
|
||||
if let Some((_, id)) = self.name_to_id.remove(name) {
|
||||
self.tokenizers.remove(&id).map(|(_, e)| e)
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
/// Clear all tokenizers from the registry
|
||||
pub fn clear(&self) {
|
||||
self.tokenizers.clear();
|
||||
self.name_to_id.clear();
|
||||
self.loading_locks.clear();
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for TokenizerRegistry {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::time::Duration;
|
||||
|
||||
use tokio::time::sleep;
|
||||
|
||||
use super::*;
|
||||
use crate::tokenizer::mock::MockTokenizer;
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_basic_operations() {
|
||||
let registry = TokenizerRegistry::new();
|
||||
|
||||
// Registry starts empty
|
||||
assert!(registry.is_empty());
|
||||
assert_eq!(registry.len(), 0);
|
||||
assert!(!registry.contains("model1"));
|
||||
|
||||
// Load and register a tokenizer
|
||||
let id = TokenizerRegistry::generate_id();
|
||||
let outcome = registry
|
||||
.load(&id, "model1", "path/to/model", || async {
|
||||
Ok(Arc::new(MockTokenizer::default()) as Arc<dyn Tokenizer>)
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Verify LoadOutcome::Loaded
|
||||
assert!(outcome.is_newly_loaded());
|
||||
assert_eq!(outcome.id(), id);
|
||||
|
||||
// Verify it's loaded
|
||||
assert!(!registry.is_empty());
|
||||
assert_eq!(registry.len(), 1);
|
||||
assert!(registry.contains("model1"));
|
||||
assert!(registry.contains_id(&id));
|
||||
|
||||
// Get returns the tokenizer
|
||||
let entry = registry.get_by_name("model1").unwrap();
|
||||
assert_eq!(entry.id, id);
|
||||
assert_eq!(entry.name, "model1");
|
||||
assert_eq!(entry.source, "path/to/model");
|
||||
|
||||
// Remove works
|
||||
let removed = registry.remove_by_id(&id);
|
||||
assert!(removed.is_some());
|
||||
assert!(registry.is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_load_returns_already_exists() {
|
||||
let registry = TokenizerRegistry::new();
|
||||
let id1 = TokenizerRegistry::generate_id();
|
||||
let id2 = TokenizerRegistry::generate_id();
|
||||
|
||||
// First load should return Loaded
|
||||
let outcome1 = registry
|
||||
.load(&id1, "model1", "source1", || async {
|
||||
Ok(Arc::new(MockTokenizer::default()) as Arc<dyn Tokenizer>)
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(outcome1.is_newly_loaded());
|
||||
assert_eq!(outcome1.id(), id1);
|
||||
|
||||
// Second load with same name should return AlreadyExists with ORIGINAL id
|
||||
let outcome2 = registry
|
||||
.load(&id2, "model1", "source2", || async {
|
||||
panic!("Loader should not be called for duplicate name");
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(!outcome2.is_newly_loaded());
|
||||
assert_eq!(outcome2.id(), id1); // Returns the original ID, not id2
|
||||
|
||||
// Registry still has only one entry
|
||||
assert_eq!(registry.len(), 1);
|
||||
|
||||
// Original source is preserved
|
||||
let entry = registry.get_by_name("model1").unwrap();
|
||||
assert_eq!(entry.source, "source1");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_load_validation() {
|
||||
let registry = TokenizerRegistry::new();
|
||||
let id = TokenizerRegistry::generate_id();
|
||||
|
||||
// Empty name should fail
|
||||
let result = registry
|
||||
.load(&id, "", "source", || async {
|
||||
panic!("Loader should not be called for invalid input");
|
||||
})
|
||||
.await;
|
||||
assert!(matches!(result, Err(LoadError::EmptyName)));
|
||||
|
||||
// Empty source should fail
|
||||
let result = registry
|
||||
.load(&id, "model", "", || async {
|
||||
panic!("Loader should not be called for invalid input");
|
||||
})
|
||||
.await;
|
||||
assert!(matches!(result, Err(LoadError::EmptySource)));
|
||||
|
||||
// Registry should be empty (nothing was loaded)
|
||||
assert!(registry.is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_load_prevents_duplicate_loading() {
|
||||
let registry = Arc::new(TokenizerRegistry::new());
|
||||
let load_count = Arc::new(std::sync::atomic::AtomicUsize::new(0));
|
||||
|
||||
// Spawn multiple tasks trying to load the same tokenizer
|
||||
let mut handles = vec![];
|
||||
for i in 0..10 {
|
||||
let registry = registry.clone();
|
||||
let load_count = load_count.clone();
|
||||
let id = format!("id-{}", i);
|
||||
let handle = tokio::spawn(async move {
|
||||
registry
|
||||
.load(&id, "model1", "source", || async {
|
||||
// Simulate slow loading
|
||||
sleep(Duration::from_millis(10)).await;
|
||||
load_count.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
|
||||
Ok(Arc::new(MockTokenizer::default()) as Arc<dyn Tokenizer>)
|
||||
})
|
||||
.await
|
||||
});
|
||||
handles.push(handle);
|
||||
}
|
||||
|
||||
// Wait for all tasks
|
||||
for handle in handles {
|
||||
handle.await.unwrap().unwrap();
|
||||
}
|
||||
|
||||
// Verify tokenizer was loaded only once
|
||||
assert_eq!(
|
||||
load_count.load(std::sync::atomic::Ordering::SeqCst),
|
||||
1,
|
||||
"Tokenizer should be loaded exactly once despite concurrent requests"
|
||||
);
|
||||
assert_eq!(registry.len(), 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_multiple_models() {
|
||||
let registry = TokenizerRegistry::new();
|
||||
|
||||
// Load multiple tokenizers
|
||||
for i in 1..=5 {
|
||||
let model_name = format!("model{}", i);
|
||||
let id = TokenizerRegistry::generate_id();
|
||||
registry
|
||||
.load(&id, &model_name, "source", || async {
|
||||
Ok(Arc::new(MockTokenizer::default()) as Arc<dyn Tokenizer>)
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
assert_eq!(registry.len(), 5);
|
||||
assert!(registry.contains("model1"));
|
||||
assert!(registry.contains("model5"));
|
||||
assert!(!registry.contains("model6"));
|
||||
|
||||
// List returns all with metadata
|
||||
let entries = registry.list();
|
||||
assert_eq!(entries.len(), 5);
|
||||
assert!(entries.iter().any(|e| e.name == "model1"));
|
||||
|
||||
// Clear all
|
||||
registry.clear();
|
||||
assert!(registry.is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_load_failure() {
|
||||
let registry = TokenizerRegistry::new();
|
||||
let id = TokenizerRegistry::generate_id();
|
||||
|
||||
// Try to load with a failing loader
|
||||
let result = registry
|
||||
.load(&id, "failing_model", "source", || async {
|
||||
Err("Load failed".to_string())
|
||||
})
|
||||
.await;
|
||||
|
||||
assert!(result.is_err());
|
||||
assert!(!registry.contains("failing_model"));
|
||||
assert!(registry.is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_get_by_name_and_id() {
|
||||
let registry = TokenizerRegistry::new();
|
||||
let id = TokenizerRegistry::generate_id();
|
||||
|
||||
registry
|
||||
.load(&id, "my-model", "hf/model", || async {
|
||||
Ok(Arc::new(MockTokenizer::default()) as Arc<dyn Tokenizer>)
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Get by name
|
||||
let by_name = registry.get_by_name("my-model");
|
||||
assert!(by_name.is_some());
|
||||
assert_eq!(by_name.as_ref().unwrap().id, id);
|
||||
|
||||
// Get by ID
|
||||
let by_id = registry.get_by_id(&id);
|
||||
assert!(by_id.is_some());
|
||||
assert_eq!(by_id.as_ref().unwrap().name, "my-model");
|
||||
|
||||
// Generic get works with both
|
||||
assert!(registry.get("my-model").is_some());
|
||||
assert!(registry.get(&id).is_some());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_register_only_if_absent() {
|
||||
let registry = TokenizerRegistry::new();
|
||||
let id1 = TokenizerRegistry::generate_id();
|
||||
let id2 = TokenizerRegistry::generate_id();
|
||||
let tokenizer1 = Arc::new(MockTokenizer::default()) as Arc<dyn Tokenizer>;
|
||||
let tokenizer2 = Arc::new(MockTokenizer::default()) as Arc<dyn Tokenizer>;
|
||||
|
||||
// First registration should succeed
|
||||
let result1 = registry.register(&id1, "model1", "source1", tokenizer1.clone());
|
||||
assert!(result1.is_some());
|
||||
assert_eq!(registry.len(), 1);
|
||||
|
||||
// Second registration with same name should fail
|
||||
let result2 = registry.register(&id2, "model1", "source2", tokenizer2.clone());
|
||||
assert!(result2.is_none());
|
||||
assert_eq!(registry.len(), 1);
|
||||
|
||||
// Original tokenizer should still be there
|
||||
let entry = registry.get_by_name("model1").unwrap();
|
||||
assert_eq!(entry.id, id1);
|
||||
assert_eq!(entry.source, "source1");
|
||||
|
||||
// Registration with different name should succeed
|
||||
let id3 = TokenizerRegistry::generate_id();
|
||||
let result3 = registry.register(&id3, "model2", "source2", tokenizer2);
|
||||
assert!(result3.is_some());
|
||||
assert_eq!(registry.len(), 2);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_loading_lock_cleanup_on_panic() {
|
||||
let registry = Arc::new(TokenizerRegistry::new());
|
||||
|
||||
// Spawn a task that will panic during loading
|
||||
let registry_clone = registry.clone();
|
||||
let handle = tokio::spawn(async move {
|
||||
registry_clone
|
||||
.load(
|
||||
&TokenizerRegistry::generate_id(),
|
||||
"panic-model",
|
||||
"source",
|
||||
|| async {
|
||||
panic!("Simulated panic during tokenizer loading");
|
||||
},
|
||||
)
|
||||
.await
|
||||
});
|
||||
|
||||
// Wait for the task - it should panic
|
||||
let result = handle.await;
|
||||
assert!(result.is_err(), "Task should have panicked");
|
||||
|
||||
// The RAII guard should have cleaned up the loading lock.
|
||||
// Verify by attempting another load with the same name - it should work,
|
||||
// not deadlock or fail due to stale lock.
|
||||
let id = TokenizerRegistry::generate_id();
|
||||
let outcome = registry
|
||||
.load(&id, "panic-model", "source", || async {
|
||||
Ok(Arc::new(MockTokenizer::default()) as Arc<dyn Tokenizer>)
|
||||
})
|
||||
.await;
|
||||
|
||||
// Should succeed - the lock was properly cleaned up
|
||||
assert!(outcome.is_ok(), "Load should succeed after panic cleanup");
|
||||
assert!(outcome.unwrap().is_newly_loaded());
|
||||
assert_eq!(registry.len(), 1);
|
||||
assert!(registry.contains("panic-model"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_loading_lock_cleanup_on_early_return() {
|
||||
let registry = Arc::new(TokenizerRegistry::new());
|
||||
|
||||
// Load a tokenizer
|
||||
let id1 = TokenizerRegistry::generate_id();
|
||||
registry
|
||||
.load(&id1, "model1", "source1", || async {
|
||||
Ok(Arc::new(MockTokenizer::default()) as Arc<dyn Tokenizer>)
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Now simulate concurrent load attempts where one thread wins
|
||||
// and another thread sees "already exists" in the double-check.
|
||||
// The RAII guard should clean up the lock on early return.
|
||||
|
||||
// First, verify loading_locks is empty after successful load
|
||||
// by checking that we can load a different model without issues
|
||||
let id2 = TokenizerRegistry::generate_id();
|
||||
let outcome = registry
|
||||
.load(&id2, "model2", "source2", || async {
|
||||
Ok(Arc::new(MockTokenizer::default()) as Arc<dyn Tokenizer>)
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert!(outcome.is_newly_loaded());
|
||||
assert_eq!(registry.len(), 2);
|
||||
|
||||
// Try to load model1 again - should return AlreadyExists
|
||||
// and the lock should be cleaned up (not leak)
|
||||
let id3 = TokenizerRegistry::generate_id();
|
||||
let outcome = registry
|
||||
.load(&id3, "model1", "source1", || async {
|
||||
panic!("Loader should not be called for existing model");
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert!(!outcome.is_newly_loaded());
|
||||
assert_eq!(outcome.id(), id1); // Returns original ID
|
||||
}
|
||||
}
|
||||
@@ -1,279 +0,0 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use anyhow::Result;
|
||||
|
||||
use super::traits::{TokenIdType, Tokenizer as TokenizerTrait};
|
||||
|
||||
/// Maintains state for an ongoing sequence of tokens and their decoded text
|
||||
/// This provides a cleaner abstraction for managing token sequences
|
||||
pub struct Sequence {
|
||||
/// The tokenizer used for encoding/decoding
|
||||
tokenizer: Arc<dyn TokenizerTrait>,
|
||||
|
||||
/// The current sequence of token ids
|
||||
token_ids: Vec<TokenIdType>,
|
||||
|
||||
/// The position in the current sequence the last decoded token completed
|
||||
prefix_offset: usize,
|
||||
|
||||
/// Current position in the sequence
|
||||
read_offset: usize,
|
||||
|
||||
/// Whether to skip special tokens when decoding
|
||||
skip_special_tokens: bool,
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for Sequence {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.debug_struct("Sequence")
|
||||
.field("tokenizer", &"Arc<dyn Tokenizer>")
|
||||
.field(
|
||||
"token_ids",
|
||||
&format_args!("{}", {
|
||||
let token_ids = self.token_ids();
|
||||
if token_ids.len() <= 20 {
|
||||
format!("{:?}", token_ids)
|
||||
} else {
|
||||
let first_ten = &token_ids[..10];
|
||||
let last_ten = &token_ids[token_ids.len() - 10..];
|
||||
format!("{:?} ... {:?}", first_ten, last_ten)
|
||||
}
|
||||
}),
|
||||
)
|
||||
.field("prefix_offset", &self.prefix_offset)
|
||||
.field("read_offset", &self.read_offset)
|
||||
.field("token count", &self.token_ids.len())
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
impl Sequence {
|
||||
/// Create a new empty sequence
|
||||
pub fn new(tokenizer: Arc<dyn TokenizerTrait>) -> Self {
|
||||
Self::new_with_options(tokenizer, false)
|
||||
}
|
||||
|
||||
/// Create a new empty sequence with skip_special_tokens option
|
||||
pub fn new_with_options(tokenizer: Arc<dyn TokenizerTrait>, skip_special_tokens: bool) -> Self {
|
||||
Self {
|
||||
tokenizer,
|
||||
token_ids: Vec::new(),
|
||||
prefix_offset: 0,
|
||||
read_offset: 0,
|
||||
skip_special_tokens,
|
||||
}
|
||||
}
|
||||
|
||||
/// Create a sequence with initial tokens
|
||||
pub fn with_tokens(tokenizer: Arc<dyn TokenizerTrait>, token_ids: Vec<TokenIdType>) -> Self {
|
||||
Self::with_tokens_and_options(tokenizer, token_ids, false)
|
||||
}
|
||||
|
||||
/// Create a sequence with initial tokens and skip_special_tokens option
|
||||
pub fn with_tokens_and_options(
|
||||
tokenizer: Arc<dyn TokenizerTrait>,
|
||||
token_ids: Vec<TokenIdType>,
|
||||
skip_special_tokens: bool,
|
||||
) -> Self {
|
||||
let len = token_ids.len();
|
||||
Self {
|
||||
tokenizer,
|
||||
token_ids,
|
||||
prefix_offset: 0,
|
||||
read_offset: len,
|
||||
skip_special_tokens,
|
||||
}
|
||||
}
|
||||
|
||||
/// Check if the sequence is empty
|
||||
#[inline]
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.token_ids.is_empty()
|
||||
}
|
||||
|
||||
/// Get the length of the sequence
|
||||
#[inline]
|
||||
pub fn len(&self) -> usize {
|
||||
self.token_ids.len()
|
||||
}
|
||||
|
||||
/// Clear the sequence
|
||||
pub fn clear(&mut self) {
|
||||
self.token_ids.clear();
|
||||
self.prefix_offset = 0;
|
||||
self.read_offset = 0;
|
||||
}
|
||||
|
||||
/// Append text to the sequence by encoding it
|
||||
///
|
||||
/// Set `add_special_tokens` to `true` for embeddings, or `false` for chat completion
|
||||
/// where the chat template already handles special tokens.
|
||||
pub fn append_text(&mut self, input: &str, add_special_tokens: bool) -> Result<()> {
|
||||
let encoding = self.tokenizer.encode(input, add_special_tokens)?;
|
||||
self.token_ids.extend(encoding.token_ids());
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Append a single token to the sequence and return newly decoded text
|
||||
/// Based on HuggingFace TGI incremental decoding
|
||||
#[inline]
|
||||
pub fn append_token(&mut self, token_id: TokenIdType) -> Result<String> {
|
||||
// Store the old read offset before adding the new token
|
||||
let old_read_offset = self.read_offset;
|
||||
|
||||
self.token_ids.push(token_id);
|
||||
self.read_offset = self.token_ids.len();
|
||||
|
||||
// If this is the first token or we're at the beginning, decode everything
|
||||
if self.prefix_offset == 0 && old_read_offset == 0 {
|
||||
let text = self
|
||||
.tokenizer
|
||||
.decode(&self.token_ids, self.skip_special_tokens)?;
|
||||
if text.ends_with("<EFBFBD>") {
|
||||
// Incomplete UTF-8 sequence, wait for more tokens
|
||||
return Ok(String::new());
|
||||
}
|
||||
self.prefix_offset = 0;
|
||||
return Ok(text);
|
||||
}
|
||||
|
||||
// Decode the text up to the previous position
|
||||
let prefix_text = self.tokenizer.decode(
|
||||
&self.token_ids[self.prefix_offset..old_read_offset],
|
||||
self.skip_special_tokens,
|
||||
)?;
|
||||
|
||||
// Decode the text including the new token
|
||||
let new_text = self.tokenizer.decode(
|
||||
&self.token_ids[self.prefix_offset..],
|
||||
self.skip_special_tokens,
|
||||
)?;
|
||||
|
||||
// Handle multi-byte character boundaries
|
||||
let mut prefix_text_len = prefix_text.len();
|
||||
while !new_text.is_char_boundary(prefix_text_len) && prefix_text_len > 0 {
|
||||
prefix_text_len -= 1;
|
||||
}
|
||||
|
||||
if new_text.len() > prefix_text.len() {
|
||||
if new_text.ends_with("<EFBFBD>") {
|
||||
// Incomplete UTF-8 sequence, wait for more tokens
|
||||
return Ok(String::new());
|
||||
} else {
|
||||
// Return the new text portion
|
||||
let incremental_text = new_text[prefix_text_len..].to_string().replace("<EFBFBD>", "");
|
||||
self.prefix_offset = old_read_offset;
|
||||
return Ok(incremental_text);
|
||||
}
|
||||
}
|
||||
|
||||
Ok(String::new())
|
||||
}
|
||||
|
||||
/// Get a reference to the tokenizer
|
||||
#[inline]
|
||||
pub fn tokenizer(&self) -> &Arc<dyn TokenizerTrait> {
|
||||
&self.tokenizer
|
||||
}
|
||||
|
||||
/// Get the current token ids
|
||||
#[inline]
|
||||
pub fn token_ids(&self) -> &[TokenIdType] {
|
||||
&self.token_ids
|
||||
}
|
||||
|
||||
/// Decode the entire sequence to text
|
||||
pub fn text(&self) -> Result<String> {
|
||||
self.tokenizer
|
||||
.decode(&self.token_ids, self.skip_special_tokens)
|
||||
}
|
||||
|
||||
/// Get the prefix offset
|
||||
#[inline]
|
||||
pub fn prefix_offset(&self) -> usize {
|
||||
self.prefix_offset
|
||||
}
|
||||
|
||||
/// Get the read offset
|
||||
#[inline]
|
||||
pub fn read_offset(&self) -> usize {
|
||||
self.read_offset
|
||||
}
|
||||
|
||||
/// Get whether special tokens are skipped during decoding
|
||||
#[inline]
|
||||
pub fn skip_special_tokens(&self) -> bool {
|
||||
self.skip_special_tokens
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::tokenizer::mock::MockTokenizer;
|
||||
|
||||
#[test]
|
||||
fn test_sequence_new() {
|
||||
let tokenizer = Arc::new(MockTokenizer::new());
|
||||
let seq = Sequence::new(tokenizer);
|
||||
assert!(seq.is_empty());
|
||||
assert_eq!(seq.len(), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_sequence_append_text() {
|
||||
let tokenizer = Arc::new(MockTokenizer::new());
|
||||
let mut seq = Sequence::new(tokenizer);
|
||||
|
||||
seq.append_text("Hello", false).unwrap();
|
||||
assert!(!seq.is_empty());
|
||||
assert!(!seq.is_empty());
|
||||
|
||||
let text = seq.text().unwrap();
|
||||
assert_eq!(text, "Hello");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_sequence_append_token() {
|
||||
let tokenizer = Arc::new(MockTokenizer::new());
|
||||
let mut seq = Sequence::new(tokenizer.clone());
|
||||
|
||||
// Start with an empty sequence and append token 1 ("Hello")
|
||||
let text1 = seq.append_token(1).unwrap();
|
||||
assert_eq!(text1, "Hello");
|
||||
|
||||
// Now append token 2 ("world")
|
||||
// The mock tokenizer will decode [1, 2] as "Hello world" (with a space)
|
||||
let text2 = seq.append_token(2).unwrap();
|
||||
// The incremental text should be " world" (with the space that the mock tokenizer adds)
|
||||
assert_eq!(text2, " world");
|
||||
|
||||
assert_eq!(seq.text().unwrap(), "Hello world");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_sequence_clear() {
|
||||
let tokenizer = Arc::new(MockTokenizer::new());
|
||||
let mut seq = Sequence::new(tokenizer);
|
||||
|
||||
seq.append_text("Hello world", false).unwrap();
|
||||
assert!(!seq.is_empty());
|
||||
|
||||
seq.clear();
|
||||
assert!(seq.is_empty());
|
||||
assert_eq!(seq.len(), 0);
|
||||
assert_eq!(seq.prefix_offset(), 0);
|
||||
assert_eq!(seq.read_offset(), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_sequence_debug() {
|
||||
let tokenizer = Arc::new(MockTokenizer::new());
|
||||
let mut seq = Sequence::new(tokenizer);
|
||||
|
||||
seq.append_text("Test", false).unwrap();
|
||||
let debug_str = format!("{:?}", seq);
|
||||
assert!(debug_str.contains("Sequence"));
|
||||
assert!(debug_str.contains("token count"));
|
||||
}
|
||||
}
|
||||
@@ -1,630 +0,0 @@
|
||||
use std::{collections::HashSet, sync::Arc};
|
||||
|
||||
use aho_corasick::AhoCorasick;
|
||||
use anyhow::Result;
|
||||
|
||||
use super::{
|
||||
sequence::Sequence,
|
||||
traits::{self, TokenIdType},
|
||||
};
|
||||
|
||||
/// Output from the sequence decoder
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub enum SequenceDecoderOutput {
|
||||
/// Normal text output
|
||||
Text(String),
|
||||
/// Text is being held due to partial stop sequence match
|
||||
Held,
|
||||
/// Stop sequence matched (hidden - not included in output)
|
||||
Stopped,
|
||||
/// Stop sequence matched with text (visible - included in output)
|
||||
StoppedWithText(String),
|
||||
}
|
||||
|
||||
/// Configuration for stop sequences
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct StopSequenceConfig {
|
||||
/// Token IDs that trigger a stop
|
||||
pub stop_tokens: HashSet<TokenIdType>,
|
||||
/// String sequences that trigger a stop
|
||||
pub stop_sequences: Vec<String>,
|
||||
/// Token IDs for visible stops (included in output)
|
||||
pub visible_stop_tokens: HashSet<TokenIdType>,
|
||||
/// String sequences for visible stops (included in output)
|
||||
pub visible_stop_sequences: Vec<String>,
|
||||
}
|
||||
|
||||
impl StopSequenceConfig {
|
||||
/// Builder pattern - add a stop token
|
||||
pub fn with_stop_token(mut self, token_id: TokenIdType) -> Self {
|
||||
self.stop_tokens.insert(token_id);
|
||||
self
|
||||
}
|
||||
|
||||
/// Builder pattern - add a stop sequence
|
||||
pub fn with_stop_sequence(mut self, sequence: impl Into<String>) -> Self {
|
||||
self.stop_sequences.push(sequence.into());
|
||||
self
|
||||
}
|
||||
|
||||
/// Builder pattern - add a visible stop token
|
||||
pub fn with_visible_stop_token(mut self, token_id: TokenIdType) -> Self {
|
||||
self.visible_stop_tokens.insert(token_id);
|
||||
self
|
||||
}
|
||||
|
||||
/// Builder pattern - add a visible stop sequence
|
||||
pub fn with_visible_stop_sequence(mut self, sequence: impl Into<String>) -> Self {
|
||||
self.visible_stop_sequences.push(sequence.into());
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
/// Decoder that handles stop sequences
|
||||
pub struct StopSequenceDecoder {
|
||||
/// Sequence for incremental decoding (replaces token_buffer + offsets)
|
||||
sequence: Sequence,
|
||||
config: StopSequenceConfig,
|
||||
aho_corasick: Option<AhoCorasick>,
|
||||
visible_boundary_idx: usize,
|
||||
/// Buffer for partial matches (the "jail")
|
||||
jail_buffer: String,
|
||||
/// Whether we've stopped
|
||||
stopped: bool,
|
||||
}
|
||||
|
||||
impl StopSequenceDecoder {
|
||||
/// Create a new stop sequence decoder
|
||||
pub fn new(
|
||||
tokenizer: Arc<dyn traits::Tokenizer>,
|
||||
config: StopSequenceConfig,
|
||||
skip_special_tokens: bool,
|
||||
) -> Self {
|
||||
let mut patterns: Vec<String> = config
|
||||
.stop_sequences
|
||||
.iter()
|
||||
.filter(|s| !s.is_empty())
|
||||
.cloned()
|
||||
.collect();
|
||||
let visible_boundary_idx = patterns.len();
|
||||
patterns.extend(
|
||||
config
|
||||
.visible_stop_sequences
|
||||
.iter()
|
||||
.filter(|s| !s.is_empty())
|
||||
.cloned(),
|
||||
);
|
||||
|
||||
let aho_corasick = if patterns.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(AhoCorasick::new(patterns).expect("Failed to build Aho-Corasick automaton"))
|
||||
};
|
||||
StopSequenceDecoder {
|
||||
sequence: Sequence::new_with_options(tokenizer, skip_special_tokens),
|
||||
config,
|
||||
aho_corasick,
|
||||
visible_boundary_idx,
|
||||
jail_buffer: String::new(),
|
||||
stopped: false,
|
||||
}
|
||||
}
|
||||
|
||||
/// Process a single token
|
||||
pub fn process_token(&mut self, token_id: TokenIdType) -> Result<SequenceDecoderOutput> {
|
||||
if self.stopped {
|
||||
return Ok(SequenceDecoderOutput::Stopped);
|
||||
}
|
||||
|
||||
// Check for token-level stops first
|
||||
if self.config.stop_tokens.contains(&token_id) {
|
||||
self.stopped = true;
|
||||
|
||||
// Flush any jailed text before stopping - use mem::take to avoid clone
|
||||
if !self.jail_buffer.is_empty() {
|
||||
return Ok(SequenceDecoderOutput::StoppedWithText(std::mem::take(
|
||||
&mut self.jail_buffer,
|
||||
)));
|
||||
}
|
||||
return Ok(SequenceDecoderOutput::Stopped);
|
||||
}
|
||||
|
||||
if self.config.visible_stop_tokens.contains(&token_id) {
|
||||
self.stopped = true;
|
||||
|
||||
// Include jailed text plus the stop token
|
||||
let stop_text = self
|
||||
.sequence
|
||||
.tokenizer()
|
||||
.decode(&[token_id], self.sequence.skip_special_tokens())?;
|
||||
let output = format!("{}{}", self.jail_buffer, stop_text);
|
||||
self.jail_buffer.clear();
|
||||
return Ok(SequenceDecoderOutput::StoppedWithText(output));
|
||||
}
|
||||
|
||||
// Use Sequence for incremental decoding
|
||||
let new_text = self.sequence.append_token(token_id)?;
|
||||
|
||||
self.jail_buffer.push_str(&new_text);
|
||||
|
||||
// Check for stop sequences
|
||||
if let Some(ac) = &self.aho_corasick {
|
||||
if let Some(mat) = ac.find(&self.jail_buffer) {
|
||||
self.stopped = true;
|
||||
let is_visible = mat.pattern().as_usize() >= self.visible_boundary_idx;
|
||||
|
||||
if is_visible {
|
||||
let output = self.jail_buffer[..mat.end()].to_string();
|
||||
self.jail_buffer.clear();
|
||||
return Ok(SequenceDecoderOutput::StoppedWithText(output));
|
||||
} else {
|
||||
let output = self.jail_buffer[..mat.start()].to_string();
|
||||
self.jail_buffer.clear();
|
||||
return Ok(if output.is_empty() {
|
||||
SequenceDecoderOutput::Stopped
|
||||
} else {
|
||||
SequenceDecoderOutput::StoppedWithText(output)
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Check for partial matches: is the end of jail_buffer the start of any stop_seq?
|
||||
// This handles stop sequences split across tokens
|
||||
let buffer_len = self.jail_buffer.len();
|
||||
let mut best_split_pos: Option<usize> = None;
|
||||
|
||||
for stop_seq in self
|
||||
.config
|
||||
.stop_sequences
|
||||
.iter()
|
||||
.chain(&self.config.visible_stop_sequences)
|
||||
{
|
||||
let stop_len = stop_seq.len();
|
||||
|
||||
if stop_len <= 1 || buffer_len == 0 {
|
||||
continue;
|
||||
}
|
||||
|
||||
let max_len = buffer_len.min(stop_len - 1);
|
||||
|
||||
for len in (1..=max_len).rev() {
|
||||
let suffix_start = buffer_len - len;
|
||||
|
||||
if !self.jail_buffer.is_char_boundary(suffix_start) {
|
||||
continue;
|
||||
}
|
||||
|
||||
let suffix = &self.jail_buffer[suffix_start..];
|
||||
|
||||
if stop_seq.starts_with(suffix)
|
||||
&& best_split_pos.is_none_or(|current| suffix_start < current)
|
||||
{
|
||||
best_split_pos = Some(suffix_start);
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(split_pos) = best_split_pos {
|
||||
// Hold the partial match, flush the rest
|
||||
// Use split_off for zero-copy: keeps [0..split_pos] in place, returns [split_pos..]
|
||||
// Then swap so we output the prefix and keep the suffix
|
||||
let suffix = self.jail_buffer.split_off(split_pos);
|
||||
let to_output = std::mem::replace(&mut self.jail_buffer, suffix);
|
||||
|
||||
if to_output.is_empty() {
|
||||
Ok(SequenceDecoderOutput::Held)
|
||||
} else {
|
||||
Ok(SequenceDecoderOutput::Text(to_output))
|
||||
}
|
||||
} else {
|
||||
// No partial matches - flush everything
|
||||
let output = std::mem::take(&mut self.jail_buffer);
|
||||
if output.is_empty() {
|
||||
Ok(SequenceDecoderOutput::Held)
|
||||
} else {
|
||||
Ok(SequenceDecoderOutput::Text(output))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Process multiple tokens
|
||||
pub fn process_tokens(
|
||||
&mut self,
|
||||
token_ids: &[TokenIdType],
|
||||
) -> Result<Vec<SequenceDecoderOutput>> {
|
||||
// Pre-allocate with exact capacity to avoid reallocations
|
||||
let mut outputs = Vec::with_capacity(token_ids.len());
|
||||
for &token_id in token_ids {
|
||||
outputs.push(self.process_token(token_id)?);
|
||||
}
|
||||
Ok(outputs)
|
||||
}
|
||||
|
||||
/// Flush any held text
|
||||
pub fn flush(&mut self) -> SequenceDecoderOutput {
|
||||
if !self.jail_buffer.is_empty() {
|
||||
// Use mem::take to avoid clone - transfers ownership and leaves empty string
|
||||
SequenceDecoderOutput::Text(std::mem::take(&mut self.jail_buffer))
|
||||
} else {
|
||||
SequenceDecoderOutput::Text(String::new())
|
||||
}
|
||||
}
|
||||
|
||||
/// Check if decoding has stopped
|
||||
pub fn is_stopped(&self) -> bool {
|
||||
self.stopped
|
||||
}
|
||||
|
||||
/// Reset the decoder state
|
||||
pub fn reset(&mut self) {
|
||||
self.jail_buffer.clear();
|
||||
self.sequence.clear();
|
||||
self.stopped = false;
|
||||
}
|
||||
}
|
||||
|
||||
/// Builder for StopSequenceDecoder
|
||||
pub struct StopSequenceDecoderBuilder {
|
||||
tokenizer: Arc<dyn traits::Tokenizer>,
|
||||
config: StopSequenceConfig,
|
||||
skip_special_tokens: bool,
|
||||
}
|
||||
|
||||
impl StopSequenceDecoderBuilder {
|
||||
pub fn new(tokenizer: Arc<dyn traits::Tokenizer>) -> Self {
|
||||
StopSequenceDecoderBuilder {
|
||||
tokenizer,
|
||||
config: StopSequenceConfig::default(),
|
||||
skip_special_tokens: true,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn stop_token(mut self, token_id: TokenIdType) -> Self {
|
||||
self.config.stop_tokens.insert(token_id);
|
||||
self
|
||||
}
|
||||
|
||||
pub fn stop_sequence(mut self, sequence: impl Into<String>) -> Self {
|
||||
self.config.stop_sequences.push(sequence.into());
|
||||
self
|
||||
}
|
||||
|
||||
pub fn visible_stop_token(mut self, token_id: TokenIdType) -> Self {
|
||||
self.config.visible_stop_tokens.insert(token_id);
|
||||
self
|
||||
}
|
||||
|
||||
pub fn visible_stop_sequence(mut self, sequence: impl Into<String>) -> Self {
|
||||
self.config.visible_stop_sequences.push(sequence.into());
|
||||
self
|
||||
}
|
||||
|
||||
pub fn skip_special_tokens(mut self, skip: bool) -> Self {
|
||||
self.skip_special_tokens = skip;
|
||||
self
|
||||
}
|
||||
|
||||
pub fn build(self) -> StopSequenceDecoder {
|
||||
StopSequenceDecoder::new(self.tokenizer, self.config, self.skip_special_tokens)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::tokenizer::mock::MockTokenizer;
|
||||
|
||||
#[test]
|
||||
fn test_stop_token_detection() {
|
||||
let tokenizer = Arc::new(MockTokenizer::new());
|
||||
let config = StopSequenceConfig::default().with_stop_token(999); // <eos> token
|
||||
|
||||
let mut decoder = StopSequenceDecoder::new(tokenizer, config, false);
|
||||
|
||||
// Process tokens before stop
|
||||
let result = decoder.process_token(1).unwrap(); // "Hello"
|
||||
assert!(matches!(result, SequenceDecoderOutput::Text(_)));
|
||||
|
||||
// Process stop token
|
||||
let result = decoder.process_token(999).unwrap(); // <eos>
|
||||
assert_eq!(result, SequenceDecoderOutput::Stopped);
|
||||
|
||||
// Further tokens should also return Stopped
|
||||
let result = decoder.process_token(2).unwrap();
|
||||
assert_eq!(result, SequenceDecoderOutput::Stopped);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_visible_stop_token() {
|
||||
let tokenizer = Arc::new(MockTokenizer::new());
|
||||
let config = StopSequenceConfig::default().with_visible_stop_token(999);
|
||||
|
||||
let mut decoder = StopSequenceDecoder::new(tokenizer, config, false);
|
||||
|
||||
let result = decoder.process_token(999).unwrap();
|
||||
assert!(matches!(result, SequenceDecoderOutput::StoppedWithText(_)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_builder_pattern() {
|
||||
let tokenizer = Arc::new(MockTokenizer::new());
|
||||
|
||||
let decoder = StopSequenceDecoderBuilder::new(tokenizer)
|
||||
.stop_token(999)
|
||||
.stop_sequence("STOP")
|
||||
.visible_stop_token(1000)
|
||||
.skip_special_tokens(true)
|
||||
.build();
|
||||
|
||||
assert!(!decoder.is_stopped());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_incremental_decoding_no_repetition() {
|
||||
// This test verifies the critical fix: no repeated output
|
||||
let tokenizer = Arc::new(MockTokenizer::new());
|
||||
let config = StopSequenceConfig::default();
|
||||
let mut decoder = StopSequenceDecoder::new(tokenizer, config, false);
|
||||
|
||||
// Process tokens one by one and collect outputs
|
||||
let mut outputs = Vec::new();
|
||||
|
||||
// Token 1: "Hello"
|
||||
let result = decoder.process_token(1).unwrap();
|
||||
if let SequenceDecoderOutput::Text(text) = result {
|
||||
outputs.push(text.clone());
|
||||
}
|
||||
|
||||
// Token 2: "world"
|
||||
let result = decoder.process_token(2).unwrap();
|
||||
if let SequenceDecoderOutput::Text(text) = result {
|
||||
outputs.push(text.clone());
|
||||
}
|
||||
|
||||
// Token 3: "test"
|
||||
let result = decoder.process_token(3).unwrap();
|
||||
if let SequenceDecoderOutput::Text(text) = result {
|
||||
outputs.push(text.clone());
|
||||
}
|
||||
|
||||
// CRITICAL: Each output should be unique (no accumulation)
|
||||
// The fix ensures we only output NEW text, not accumulated text
|
||||
assert_eq!(outputs.len(), 3);
|
||||
|
||||
for i in 0..outputs.len() {
|
||||
for j in i + 1..outputs.len() {
|
||||
// No output should contain another (no accumulation)
|
||||
assert!(!outputs[j].contains(&outputs[i]));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_stop_sequence_detection() {
|
||||
let tokenizer = Arc::new(MockTokenizer::new());
|
||||
let config = StopSequenceConfig::default().with_stop_sequence("test");
|
||||
let mut decoder = StopSequenceDecoder::new(tokenizer, config, false);
|
||||
|
||||
// Process "Hello world"
|
||||
decoder.process_token(1).unwrap(); // "Hello"
|
||||
decoder.process_token(2).unwrap(); // "world"
|
||||
|
||||
// Process "test" which should trigger stop
|
||||
let result = decoder.process_token(3).unwrap(); // "test"
|
||||
|
||||
// Should stop when we hit "test"
|
||||
assert!(matches!(
|
||||
result,
|
||||
SequenceDecoderOutput::Stopped | SequenceDecoderOutput::StoppedWithText(_)
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_flush_after_partial() {
|
||||
let tokenizer = Arc::new(MockTokenizer::new());
|
||||
let config = StopSequenceConfig::default().with_stop_sequence("NEVER_MATCH");
|
||||
let mut decoder = StopSequenceDecoder::new(tokenizer, config, false);
|
||||
|
||||
// Process a token
|
||||
decoder.process_token(1).unwrap(); // "Hello"
|
||||
|
||||
// Flush should return any remaining text in jail
|
||||
let result = decoder.flush();
|
||||
|
||||
// After processing, flush should work
|
||||
assert!(matches!(result, SequenceDecoderOutput::Text(_)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_reset_functionality() {
|
||||
let tokenizer = Arc::new(MockTokenizer::new());
|
||||
let config = StopSequenceConfig::default().with_stop_token(999);
|
||||
let mut decoder = StopSequenceDecoder::new(tokenizer, config, false);
|
||||
|
||||
// Process and stop
|
||||
decoder.process_token(1).unwrap();
|
||||
decoder.process_token(999).unwrap();
|
||||
assert!(decoder.is_stopped());
|
||||
|
||||
// Reset should clear everything
|
||||
decoder.reset();
|
||||
assert!(!decoder.is_stopped());
|
||||
|
||||
// Should be able to process again
|
||||
let result = decoder.process_token(2).unwrap();
|
||||
assert!(matches!(result, SequenceDecoderOutput::Text(_)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_visible_stop_sequence() {
|
||||
let tokenizer = Arc::new(MockTokenizer::new());
|
||||
let config = StopSequenceConfig::default().with_visible_stop_sequence("world");
|
||||
let mut decoder = StopSequenceDecoder::new(tokenizer, config, false);
|
||||
|
||||
// Process "Hello"
|
||||
decoder.process_token(1).unwrap();
|
||||
|
||||
// Process "world" - should include it in output
|
||||
let result = decoder.process_token(2).unwrap();
|
||||
|
||||
if let SequenceDecoderOutput::StoppedWithText(text) = result {
|
||||
// Should include "world" in the output
|
||||
assert!(text.contains("world"));
|
||||
} else {
|
||||
panic!("Expected StoppedWithText with visible stop sequence");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_multiple_tokens_processing() {
|
||||
let tokenizer = Arc::new(MockTokenizer::new());
|
||||
let config = StopSequenceConfig::default();
|
||||
let mut decoder = StopSequenceDecoder::new(tokenizer, config, false);
|
||||
|
||||
// Process multiple tokens at once
|
||||
let results = decoder.process_tokens(&[1, 2, 3]).unwrap();
|
||||
|
||||
// Should get results for each token
|
||||
assert_eq!(results.len(), 3);
|
||||
|
||||
// Each result should be Text (no stops configured)
|
||||
for result in results {
|
||||
assert!(matches!(
|
||||
result,
|
||||
SequenceDecoderOutput::Text(_) | SequenceDecoderOutput::Held
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_utf8_multibyte_character_boundaries() {
|
||||
// This test verifies the fix for the UTF-8 boundary panic
|
||||
// The panic occurred when trying to slice jail_buffer at a byte index
|
||||
// that was in the middle of a multi-byte UTF-8 character (e.g., '×')
|
||||
use crate::tokenizer::mock::MockTokenizer;
|
||||
|
||||
let tokenizer = Arc::new(MockTokenizer::new());
|
||||
|
||||
// Configure stop sequence with a multi-byte character
|
||||
let config = StopSequenceConfig::default().with_stop_sequence(" ×");
|
||||
|
||||
let mut decoder = StopSequenceDecoder::new(tokenizer, config, false);
|
||||
|
||||
// Simulate the scenario: jail_buffer will contain " ×" (space + multiplication sign)
|
||||
// The '×' character is UTF-8 encoded as bytes [0xC3, 0x97] (2 bytes)
|
||||
// When checking for partial matches, we must not slice in the middle of these bytes
|
||||
|
||||
// This should not panic - the fix ensures we only slice at char boundaries
|
||||
let result = decoder.process_token(1); // Will add some text to jail_buffer
|
||||
assert!(result.is_ok());
|
||||
|
||||
// Even with multi-byte UTF-8 characters in the buffer, processing should work
|
||||
let result = decoder.process_token(2);
|
||||
assert!(result.is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_utf8_multibyte_delta_character() {
|
||||
// Test for: byte index 1 is not a char boundary; it is inside 'Δ' (bytes 0..2) of `Δ`
|
||||
// 'Δ' (U+0394 GREEK CAPITAL LETTER DELTA) is encoded as [0xCE, 0x94] (2 bytes)
|
||||
let tokenizer = Arc::new(MockTokenizer::new());
|
||||
let config = StopSequenceConfig::default().with_stop_sequence("Δ");
|
||||
|
||||
let mut decoder = StopSequenceDecoder::new(tokenizer, config, false);
|
||||
|
||||
// Process tokens - should not panic when checking partial matches
|
||||
let result = decoder.process_token(1);
|
||||
assert!(result.is_ok());
|
||||
let result = decoder.process_token(2);
|
||||
assert!(result.is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_utf8_multibyte_degree_character() {
|
||||
// Test for: byte index 1 is not a char boundary; it is inside '°' (bytes 0..2) of `°`
|
||||
// '°' (U+00B0 DEGREE SIGN) is encoded as [0xC2, 0xB0] (2 bytes)
|
||||
let tokenizer = Arc::new(MockTokenizer::new());
|
||||
let config = StopSequenceConfig::default().with_stop_sequence("°");
|
||||
|
||||
let mut decoder = StopSequenceDecoder::new(tokenizer, config, false);
|
||||
|
||||
// Process tokens - should not panic when checking partial matches
|
||||
let result = decoder.process_token(1);
|
||||
assert!(result.is_ok());
|
||||
let result = decoder.process_token(2);
|
||||
assert!(result.is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_utf8_multibyte_triangle_character() {
|
||||
// Test for: byte index 4 is not a char boundary; it is inside '∆' (bytes 2..5) of ` (∆`
|
||||
// '∆' (U+2206 INCREMENT) is encoded as [0xE2, 0x88, 0x86] (3 bytes)
|
||||
let tokenizer = Arc::new(MockTokenizer::new());
|
||||
let config = StopSequenceConfig::default().with_stop_sequence(" (∆");
|
||||
|
||||
let mut decoder = StopSequenceDecoder::new(tokenizer, config, false);
|
||||
|
||||
// Process tokens - should not panic when checking partial matches
|
||||
let result = decoder.process_token(1);
|
||||
assert!(result.is_ok());
|
||||
let result = decoder.process_token(2);
|
||||
assert!(result.is_ok());
|
||||
let result = decoder.process_token(3);
|
||||
assert!(result.is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_utf8_multibyte_en_dash_character() {
|
||||
// Test for: byte index 3 is not a char boundary; it is inside '–' (bytes 1..4) of ` –`
|
||||
// '–' (U+2013 EN DASH) is encoded as [0xE2, 0x80, 0x93] (3 bytes)
|
||||
let tokenizer = Arc::new(MockTokenizer::new());
|
||||
let config = StopSequenceConfig::default().with_stop_sequence(" –");
|
||||
|
||||
let mut decoder = StopSequenceDecoder::new(tokenizer, config, false);
|
||||
|
||||
// Process tokens - should not panic when checking partial matches
|
||||
let result = decoder.process_token(1);
|
||||
assert!(result.is_ok());
|
||||
let result = decoder.process_token(2);
|
||||
assert!(result.is_ok());
|
||||
let result = decoder.process_token(3);
|
||||
assert!(result.is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_utf8_multibyte_various_characters() {
|
||||
// Comprehensive test with multiple multi-byte UTF-8 characters
|
||||
// Tests 2-byte, 3-byte, and 4-byte UTF-8 sequences
|
||||
let test_cases = vec![
|
||||
("×", "multiplication sign - 2 bytes"),
|
||||
("Δ", "Greek Delta - 2 bytes"),
|
||||
("°", "degree sign - 2 bytes"),
|
||||
("∆", "increment - 3 bytes"),
|
||||
("–", "en dash - 3 bytes"),
|
||||
("€", "euro sign - 3 bytes"),
|
||||
("中", "Chinese character - 3 bytes"),
|
||||
("🚀", "rocket emoji - 4 bytes"),
|
||||
("💡", "lightbulb emoji - 4 bytes"),
|
||||
];
|
||||
|
||||
for (stop_char, description) in test_cases {
|
||||
let tokenizer = Arc::new(MockTokenizer::new());
|
||||
let config = StopSequenceConfig::default().with_stop_sequence(stop_char);
|
||||
|
||||
let mut decoder = StopSequenceDecoder::new(tokenizer, config, false);
|
||||
|
||||
// Process multiple tokens - should not panic
|
||||
for token_id in 1..=5 {
|
||||
let result = decoder.process_token(token_id);
|
||||
assert!(
|
||||
result.is_ok(),
|
||||
"Failed on {} with token {}",
|
||||
description,
|
||||
token_id
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,109 +0,0 @@
|
||||
// src/tokenizer/stream.rs
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use anyhow::Result;
|
||||
|
||||
use super::traits::{self, TokenIdType};
|
||||
|
||||
const INITIAL_INCREMENTAL_DETOKENIZATION_OFFSET: usize = 5;
|
||||
|
||||
/// DecodeStream will keep the state necessary to produce individual chunks of
|
||||
/// strings given an input stream of token_ids
|
||||
pub struct DecodeStream {
|
||||
/// The tokenizer used to decode token_ids
|
||||
tokenizer: Arc<dyn traits::Tokenizer>,
|
||||
|
||||
skip_special_tokens: bool,
|
||||
|
||||
/// A temporary buffer of the necessary token_ids needed
|
||||
/// to produce valid string chunks
|
||||
all_token_ids: Vec<TokenIdType>,
|
||||
|
||||
prefix_offset: usize,
|
||||
read_offset: usize,
|
||||
}
|
||||
|
||||
impl DecodeStream {
|
||||
pub fn new(
|
||||
tokenizer: Arc<dyn traits::Tokenizer>,
|
||||
prompt_token_ids: &[TokenIdType],
|
||||
skip_special_tokens: bool,
|
||||
) -> Self {
|
||||
let num_input_tokens = prompt_token_ids.len();
|
||||
let prompt_token_ids = prompt_token_ids.to_vec();
|
||||
Self {
|
||||
tokenizer,
|
||||
skip_special_tokens,
|
||||
all_token_ids: prompt_token_ids,
|
||||
prefix_offset: num_input_tokens
|
||||
.saturating_sub(INITIAL_INCREMENTAL_DETOKENIZATION_OFFSET),
|
||||
read_offset: num_input_tokens,
|
||||
}
|
||||
}
|
||||
|
||||
/// Step appends a token_id to the internal state and tries to produce a text chunk.
|
||||
/// Returning `None` means the given id is not enough to produce a chunk.
|
||||
#[inline]
|
||||
pub fn step(&mut self, id: TokenIdType) -> Result<Option<String>> {
|
||||
self.all_token_ids.push(id);
|
||||
|
||||
let prefix_text = self.tokenizer.decode(
|
||||
&self.all_token_ids[self.prefix_offset..self.read_offset],
|
||||
self.skip_special_tokens,
|
||||
)?;
|
||||
|
||||
let new_text = self.tokenizer.decode(
|
||||
&self.all_token_ids[self.prefix_offset..],
|
||||
self.skip_special_tokens,
|
||||
)?;
|
||||
|
||||
if new_text.len() > prefix_text.len() && !new_text.ends_with("<EFBFBD>") {
|
||||
let new_text = new_text[prefix_text.len()..].to_string();
|
||||
|
||||
self.prefix_offset = self.read_offset;
|
||||
self.read_offset = self.all_token_ids.len();
|
||||
|
||||
Ok(Some(new_text))
|
||||
} else {
|
||||
Ok(None)
|
||||
}
|
||||
}
|
||||
|
||||
/// Process multiple tokens at once
|
||||
pub fn step_batch(&mut self, token_ids: &[u32]) -> Result<Vec<String>> {
|
||||
// Pre-allocate with capacity - most tokens produce output
|
||||
let mut chunks = Vec::with_capacity(token_ids.len());
|
||||
|
||||
for &token_id in token_ids {
|
||||
if let Some(text) = self.step(token_id)? {
|
||||
chunks.push(text);
|
||||
}
|
||||
}
|
||||
|
||||
Ok(chunks)
|
||||
}
|
||||
|
||||
/// Force flush any remaining text
|
||||
pub fn flush(&mut self) -> Result<Option<String>> {
|
||||
if self.read_offset < self.all_token_ids.len() {
|
||||
let remaining = self.tokenizer.decode(
|
||||
&self.all_token_ids[self.read_offset..],
|
||||
self.skip_special_tokens,
|
||||
)?;
|
||||
|
||||
self.read_offset = self.all_token_ids.len();
|
||||
|
||||
if !remaining.is_empty() {
|
||||
return Ok(Some(remaining));
|
||||
}
|
||||
}
|
||||
|
||||
Ok(None)
|
||||
}
|
||||
|
||||
/// Get all tokens processed so far
|
||||
pub fn tokens(&self) -> &[u32] {
|
||||
&self.all_token_ids
|
||||
}
|
||||
}
|
||||
@@ -1,139 +0,0 @@
|
||||
#[cfg(test)]
|
||||
use std::sync::Arc;
|
||||
|
||||
#[cfg(test)]
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_mock_tokenizer_encode() {
|
||||
let tokenizer = mock::MockTokenizer::new();
|
||||
let encoding = tokenizer.encode("Hello world", false).unwrap();
|
||||
let token_ids = encoding.token_ids();
|
||||
assert_eq!(token_ids, &[1, 2]); // "Hello" -> 1, "world" -> 2
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_mock_tokenizer_decode() {
|
||||
let tokenizer = mock::MockTokenizer::new();
|
||||
let text = tokenizer.decode(&[1, 2], false).unwrap();
|
||||
assert_eq!(text, "Hello world");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_mock_tokenizer_decode_skip_special() {
|
||||
let tokenizer = mock::MockTokenizer::new();
|
||||
|
||||
// With special tokens
|
||||
let text = tokenizer.decode(&[1000, 1, 2, 999], false).unwrap();
|
||||
assert_eq!(text, "<bos> Hello world <eos>");
|
||||
|
||||
// Without special tokens
|
||||
let text = tokenizer.decode(&[1000, 1, 2, 999], true).unwrap();
|
||||
assert_eq!(text, "Hello world");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_tokenizer_wrapper() {
|
||||
let mock_tokenizer = Arc::new(mock::MockTokenizer::new());
|
||||
let tokenizer = Tokenizer::from_arc(mock_tokenizer);
|
||||
|
||||
let encoding = tokenizer.encode("Hello world", false).unwrap();
|
||||
assert_eq!(encoding.token_ids(), &[1, 2]);
|
||||
|
||||
let text = tokenizer.decode(&[1, 2], false).unwrap();
|
||||
assert_eq!(text, "Hello world");
|
||||
|
||||
assert_eq!(tokenizer.vocab_size(), 14);
|
||||
|
||||
assert_eq!(tokenizer.token_to_id("Hello"), Some(1));
|
||||
assert_eq!(tokenizer.token_to_id("unknown"), None);
|
||||
|
||||
assert_eq!(tokenizer.id_to_token(1), Some("Hello".to_string()));
|
||||
assert_eq!(tokenizer.id_to_token(9999), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_decode_stream_basic() {
|
||||
let mock_tokenizer = Arc::new(mock::MockTokenizer::new());
|
||||
let tokenizer = Tokenizer::from_arc(mock_tokenizer);
|
||||
|
||||
// Create a decode stream with initial tokens
|
||||
let initial_tokens = vec![1, 2]; // "Hello world"
|
||||
let mut stream = tokenizer.decode_stream(&initial_tokens, false);
|
||||
|
||||
// Add a new token
|
||||
let result = stream.step(3).unwrap(); // "test"
|
||||
// Since we're using a mock, the actual incremental behavior depends on implementation
|
||||
// For now, we just verify it doesn't crash
|
||||
assert!(result.is_some() || result.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_decode_stream_flush() {
|
||||
let mock_tokenizer = Arc::new(mock::MockTokenizer::new());
|
||||
let tokenizer = Tokenizer::from_arc(mock_tokenizer);
|
||||
|
||||
let initial_tokens = vec![1];
|
||||
let mut stream = tokenizer.decode_stream(&initial_tokens, false);
|
||||
|
||||
// Add tokens
|
||||
stream.step(2).unwrap();
|
||||
stream.step(3).unwrap();
|
||||
|
||||
// Flush remaining
|
||||
let flushed = stream.flush().unwrap();
|
||||
// The flush behavior depends on the implementation
|
||||
assert!(flushed.is_some() || flushed.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_special_tokens() {
|
||||
let mock_tokenizer = Arc::new(mock::MockTokenizer::new());
|
||||
let tokenizer = Tokenizer::from_arc(mock_tokenizer);
|
||||
|
||||
let special_tokens = tokenizer.get_special_tokens();
|
||||
assert_eq!(special_tokens.bos_token, Some("<bos>".to_string()));
|
||||
assert_eq!(special_tokens.eos_token, Some("<eos>".to_string()));
|
||||
assert_eq!(special_tokens.unk_token, Some("<unk>".to_string()));
|
||||
assert!(special_tokens.sep_token.is_none());
|
||||
assert!(special_tokens.pad_token.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_batch_encode() {
|
||||
let tokenizer = mock::MockTokenizer::new();
|
||||
let inputs = vec!["Hello", "world", "test"];
|
||||
let encodings = tokenizer.encode_batch(&inputs, false).unwrap();
|
||||
|
||||
assert_eq!(encodings.len(), 3);
|
||||
assert_eq!(encodings[0].token_ids(), &[1]); // "Hello" -> 1
|
||||
assert_eq!(encodings[1].token_ids(), &[2]); // "world" -> 2
|
||||
assert_eq!(encodings[2].token_ids(), &[3]); // "test" -> 3
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_thread_safety() {
|
||||
use std::thread;
|
||||
|
||||
let mock_tokenizer = Arc::new(mock::MockTokenizer::new());
|
||||
let tokenizer = Tokenizer::from_arc(mock_tokenizer);
|
||||
|
||||
// Spawn multiple threads that use the same tokenizer
|
||||
let handles: Vec<_> = (0..10)
|
||||
.map(|i| {
|
||||
let tokenizer_clone = tokenizer.clone();
|
||||
thread::spawn(move || {
|
||||
let text = "Hello test".to_string();
|
||||
let encoding = tokenizer_clone.encode(&text, false).unwrap();
|
||||
let decoded = tokenizer_clone.decode(encoding.token_ids(), false).unwrap();
|
||||
assert!(decoded.contains("Hello") || decoded.contains("test"));
|
||||
i
|
||||
})
|
||||
})
|
||||
.collect();
|
||||
|
||||
// Wait for all threads to complete
|
||||
for handle in handles {
|
||||
handle.join().unwrap();
|
||||
}
|
||||
}
|
||||
@@ -1,283 +0,0 @@
|
||||
use anyhow::{Error, Result};
|
||||
use tiktoken_rs::{cl100k_base, p50k_base, p50k_edit, r50k_base, CoreBPE};
|
||||
|
||||
use super::traits::{
|
||||
Decoder, Encoder, Encoding, SpecialTokens, TokenIdType, Tokenizer as TokenizerTrait,
|
||||
};
|
||||
|
||||
/// Tiktoken tokenizer wrapper for OpenAI GPT models
|
||||
pub(crate) struct TiktokenTokenizer {
|
||||
tokenizer: CoreBPE,
|
||||
#[allow(dead_code)]
|
||||
model: TiktokenModel,
|
||||
special_tokens: SpecialTokens,
|
||||
vocab_size: usize,
|
||||
}
|
||||
|
||||
/// Supported Tiktoken models
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
pub enum TiktokenModel {
|
||||
/// GPT-4, GPT-3.5-turbo, text-embedding-ada-002
|
||||
Cl100kBase,
|
||||
/// Codex models, text-davinci-002, text-davinci-003
|
||||
P50kBase,
|
||||
/// Use for edit models like text-davinci-edit-001, code-davinci-edit-001
|
||||
P50kEdit,
|
||||
/// GPT-3 models like davinci
|
||||
R50kBase,
|
||||
}
|
||||
|
||||
impl TiktokenTokenizer {
|
||||
/// Create a new Tiktoken tokenizer for the specified model
|
||||
pub fn new(model: TiktokenModel) -> Result<Self> {
|
||||
let tokenizer =
|
||||
match model {
|
||||
TiktokenModel::Cl100kBase => cl100k_base()
|
||||
.map_err(|e| Error::msg(format!("Failed to load cl100k_base: {}", e)))?,
|
||||
TiktokenModel::P50kBase => p50k_base()
|
||||
.map_err(|e| Error::msg(format!("Failed to load p50k_base: {}", e)))?,
|
||||
TiktokenModel::P50kEdit => p50k_edit()
|
||||
.map_err(|e| Error::msg(format!("Failed to load p50k_edit: {}", e)))?,
|
||||
TiktokenModel::R50kBase => r50k_base()
|
||||
.map_err(|e| Error::msg(format!("Failed to load r50k_base: {}", e)))?,
|
||||
};
|
||||
|
||||
// Extract special tokens (tiktoken-rs doesn't expose them directly)
|
||||
// We'll use common ones for GPT models
|
||||
let special_tokens = Self::get_special_tokens_for_model(model);
|
||||
|
||||
// Get vocabulary size (this is an approximation)
|
||||
let vocab_size = match model {
|
||||
TiktokenModel::Cl100kBase => 100256, // cl100k has ~100k tokens
|
||||
TiktokenModel::P50kBase | TiktokenModel::P50kEdit => 50281, // p50k has ~50k tokens
|
||||
TiktokenModel::R50kBase => 50257, // r50k has ~50k tokens
|
||||
};
|
||||
|
||||
Ok(TiktokenTokenizer {
|
||||
tokenizer,
|
||||
model,
|
||||
special_tokens,
|
||||
vocab_size,
|
||||
})
|
||||
}
|
||||
|
||||
/// Create a tokenizer from a model string (e.g., "gpt-4", "gpt-3.5-turbo")
|
||||
pub fn from_model_name(model_name: &str) -> Result<Self> {
|
||||
let model = Self::model_from_name(model_name)?;
|
||||
Self::new(model)
|
||||
}
|
||||
|
||||
/// Determine the appropriate model from a model name
|
||||
fn model_from_name(model_name: &str) -> Result<TiktokenModel> {
|
||||
// Based on OpenAI's model-to-encoding mapping
|
||||
if model_name.contains("gpt-4")
|
||||
|| model_name.contains("gpt-3.5")
|
||||
|| model_name.contains("turbo")
|
||||
{
|
||||
Ok(TiktokenModel::Cl100kBase)
|
||||
} else if model_name.contains("davinci-002")
|
||||
|| model_name.contains("davinci-003")
|
||||
|| model_name.contains("codex")
|
||||
{
|
||||
Ok(TiktokenModel::P50kBase)
|
||||
} else if model_name.contains("edit") {
|
||||
Ok(TiktokenModel::P50kEdit)
|
||||
} else if model_name.contains("davinci")
|
||||
|| model_name.contains("curie")
|
||||
|| model_name.contains("babbage")
|
||||
|| model_name.contains("ada")
|
||||
{
|
||||
Ok(TiktokenModel::R50kBase)
|
||||
} else {
|
||||
// Return an error for unrecognized model names to prevent silent failures
|
||||
Err(anyhow::anyhow!(
|
||||
"Unrecognized OpenAI model name: '{}'. Expected GPT-3, GPT-3.5, GPT-4, or related model names",
|
||||
model_name
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
/// Get special tokens for a specific model
|
||||
fn get_special_tokens_for_model(model: TiktokenModel) -> SpecialTokens {
|
||||
// These are common special tokens for GPT models
|
||||
// The actual token IDs might vary by model
|
||||
match model {
|
||||
TiktokenModel::Cl100kBase => SpecialTokens {
|
||||
bos_token: Some("<|endoftext|>".to_string()),
|
||||
eos_token: Some("<|endoftext|>".to_string()),
|
||||
unk_token: None,
|
||||
sep_token: None,
|
||||
pad_token: Some("<|endoftext|>".to_string()),
|
||||
cls_token: None,
|
||||
mask_token: None,
|
||||
additional_special_tokens: vec![
|
||||
"<|fim_prefix|>".to_string(),
|
||||
"<|fim_middle|>".to_string(),
|
||||
"<|fim_suffix|>".to_string(),
|
||||
"<|endofprompt|>".to_string(),
|
||||
],
|
||||
},
|
||||
_ => SpecialTokens {
|
||||
bos_token: Some("<|endoftext|>".to_string()),
|
||||
eos_token: Some("<|endoftext|>".to_string()),
|
||||
unk_token: None,
|
||||
sep_token: None,
|
||||
pad_token: Some("<|endoftext|>".to_string()),
|
||||
cls_token: None,
|
||||
mask_token: None,
|
||||
additional_special_tokens: vec![],
|
||||
},
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Encoder for TiktokenTokenizer {
|
||||
fn encode(&self, input: &str, _add_special_tokens: bool) -> Result<Encoding> {
|
||||
// tiktoken uses encode_ordinary which doesn't add special tokens
|
||||
let tokens = self.tokenizer.encode_ordinary(input);
|
||||
Ok(Encoding::Tiktoken(tokens))
|
||||
}
|
||||
|
||||
fn encode_batch(&self, inputs: &[&str], add_special_tokens: bool) -> Result<Vec<Encoding>> {
|
||||
inputs
|
||||
.iter()
|
||||
.map(|input| self.encode(input, add_special_tokens))
|
||||
.collect()
|
||||
}
|
||||
}
|
||||
|
||||
impl Decoder for TiktokenTokenizer {
|
||||
fn decode(&self, token_ids: &[TokenIdType], _skip_special_tokens: bool) -> Result<String> {
|
||||
// tiktoken-rs 0.7.0 now uses u32 (Rank type)
|
||||
self.tokenizer
|
||||
.decode(token_ids.to_vec())
|
||||
.map_err(|e| Error::msg(format!("Decoding failed: {}", e)))
|
||||
}
|
||||
}
|
||||
|
||||
impl TokenizerTrait for TiktokenTokenizer {
|
||||
fn vocab_size(&self) -> usize {
|
||||
self.vocab_size
|
||||
}
|
||||
|
||||
fn get_special_tokens(&self) -> &SpecialTokens {
|
||||
&self.special_tokens
|
||||
}
|
||||
|
||||
fn token_to_id(&self, _token: &str) -> Option<TokenIdType> {
|
||||
// Tiktoken doesn't provide direct token-to-id mapping
|
||||
// We'd need to encode the token and check if it produces a single ID
|
||||
None
|
||||
}
|
||||
|
||||
fn id_to_token(&self, _id: TokenIdType) -> Option<String> {
|
||||
// Tiktoken doesn't provide direct id-to-token mapping
|
||||
// We can only decode IDs to text
|
||||
None
|
||||
}
|
||||
|
||||
fn as_any(&self) -> &dyn std::any::Any {
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_tiktoken_creation() {
|
||||
let tokenizer = TiktokenTokenizer::new(TiktokenModel::Cl100kBase).unwrap();
|
||||
assert_eq!(tokenizer.vocab_size(), 100256);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_model_from_name() {
|
||||
assert!(matches!(
|
||||
TiktokenTokenizer::model_from_name("gpt-4").unwrap(),
|
||||
TiktokenModel::Cl100kBase
|
||||
));
|
||||
assert!(matches!(
|
||||
TiktokenTokenizer::model_from_name("gpt-3.5-turbo").unwrap(),
|
||||
TiktokenModel::Cl100kBase
|
||||
));
|
||||
assert!(matches!(
|
||||
TiktokenTokenizer::model_from_name("text-davinci-003").unwrap(),
|
||||
TiktokenModel::P50kBase
|
||||
));
|
||||
assert!(matches!(
|
||||
TiktokenTokenizer::model_from_name("text-davinci-edit-001").unwrap(),
|
||||
TiktokenModel::P50kEdit
|
||||
));
|
||||
assert!(matches!(
|
||||
TiktokenTokenizer::model_from_name("davinci").unwrap(),
|
||||
TiktokenModel::R50kBase
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_encode_decode() {
|
||||
let tokenizer = TiktokenTokenizer::new(TiktokenModel::Cl100kBase).unwrap();
|
||||
|
||||
let text = "Hello, world!";
|
||||
let encoding = tokenizer.encode(text, false).unwrap();
|
||||
|
||||
let decoded = tokenizer.decode(encoding.token_ids(), false).unwrap();
|
||||
assert_eq!(decoded, text);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_batch_encode() {
|
||||
let tokenizer = TiktokenTokenizer::new(TiktokenModel::Cl100kBase).unwrap();
|
||||
|
||||
let texts = vec!["Hello", "World", "Test"];
|
||||
let encodings = tokenizer.encode_batch(&texts, false).unwrap();
|
||||
|
||||
assert_eq!(encodings.len(), 3);
|
||||
for (i, encoding) in encodings.iter().enumerate() {
|
||||
let decoded = tokenizer.decode(encoding.token_ids(), false).unwrap();
|
||||
assert_eq!(decoded, texts[i]);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_special_tokens() {
|
||||
let tokenizer = TiktokenTokenizer::new(TiktokenModel::Cl100kBase).unwrap();
|
||||
let special_tokens = tokenizer.get_special_tokens();
|
||||
|
||||
assert!(special_tokens.eos_token.is_some());
|
||||
assert_eq!(special_tokens.eos_token.as_ref().unwrap(), "<|endoftext|>");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_unrecognized_model_name_returns_error() {
|
||||
let result = TiktokenTokenizer::from_model_name("distilgpt-2");
|
||||
assert!(result.is_err());
|
||||
if let Err(e) = result {
|
||||
assert!(e.to_string().contains("Unrecognized OpenAI model name"));
|
||||
}
|
||||
|
||||
let result = TiktokenTokenizer::from_model_name("bert-base-uncased");
|
||||
assert!(result.is_err());
|
||||
if let Err(e) = result {
|
||||
assert!(e.to_string().contains("Unrecognized OpenAI model name"));
|
||||
}
|
||||
|
||||
let result = TiktokenTokenizer::from_model_name("llama-7b");
|
||||
assert!(result.is_err());
|
||||
if let Err(e) = result {
|
||||
assert!(e.to_string().contains("Unrecognized OpenAI model name"));
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_recognized_model_names() {
|
||||
assert!(TiktokenTokenizer::from_model_name("gpt-4").is_ok());
|
||||
assert!(TiktokenTokenizer::from_model_name("gpt-3.5-turbo").is_ok());
|
||||
assert!(TiktokenTokenizer::from_model_name("text-davinci-003").is_ok());
|
||||
assert!(TiktokenTokenizer::from_model_name("code-davinci-002").is_ok());
|
||||
assert!(TiktokenTokenizer::from_model_name("text-curie-001").is_ok());
|
||||
assert!(TiktokenTokenizer::from_model_name("text-babbage-001").is_ok());
|
||||
assert!(TiktokenTokenizer::from_model_name("text-ada-001").is_ok());
|
||||
}
|
||||
}
|
||||
@@ -1,90 +0,0 @@
|
||||
use std::{
|
||||
collections::hash_map::DefaultHasher,
|
||||
hash::{Hash, Hasher},
|
||||
};
|
||||
|
||||
use anyhow::Result;
|
||||
|
||||
/// Type alias for token IDs
|
||||
pub type TokenIdType = u32;
|
||||
|
||||
/// Core encoding trait - separate from decoding for modularity
|
||||
pub trait Encoder: Send + Sync {
|
||||
fn encode(&self, input: &str, add_special_tokens: bool) -> Result<Encoding>;
|
||||
fn encode_batch(&self, inputs: &[&str], add_special_tokens: bool) -> Result<Vec<Encoding>>;
|
||||
}
|
||||
|
||||
/// Core decoding trait - can be implemented independently
|
||||
pub trait Decoder: Send + Sync {
|
||||
fn decode(&self, token_ids: &[TokenIdType], skip_special_tokens: bool) -> Result<String>;
|
||||
}
|
||||
|
||||
/// Combined tokenizer trait
|
||||
pub trait Tokenizer: Encoder + Decoder {
|
||||
fn vocab_size(&self) -> usize;
|
||||
fn get_special_tokens(&self) -> &SpecialTokens;
|
||||
fn token_to_id(&self, token: &str) -> Option<TokenIdType>;
|
||||
fn id_to_token(&self, id: TokenIdType) -> Option<String>;
|
||||
|
||||
/// Enable downcasting to concrete types
|
||||
fn as_any(&self) -> &dyn std::any::Any;
|
||||
}
|
||||
|
||||
/// Contains the results of tokenizing text: token IDs, string tokens, and their spans
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum Encoding {
|
||||
/// Hugging Face
|
||||
Hf(Box<tokenizers::tokenizer::Encoding>),
|
||||
/// Sentence Piece
|
||||
Sp(Vec<TokenIdType>),
|
||||
/// Tiktoken (for GPT models) - now uses u32 in tiktoken-rs 0.7.0
|
||||
Tiktoken(Vec<TokenIdType>),
|
||||
}
|
||||
|
||||
impl Encoding {
|
||||
/// Returns a reference to token IDs - zero-copy operation
|
||||
#[inline]
|
||||
pub fn token_ids(&self) -> &[TokenIdType] {
|
||||
match self {
|
||||
Encoding::Hf(inner) => inner.get_ids(),
|
||||
Encoding::Sp(inner) => inner,
|
||||
Encoding::Tiktoken(inner) => inner,
|
||||
}
|
||||
}
|
||||
|
||||
/// Deprecated: Use token_ids() instead (kept for compatibility)
|
||||
#[deprecated(since = "0.1.0", note = "Use token_ids() instead")]
|
||||
pub fn token_ids_ref(&self) -> &[TokenIdType] {
|
||||
self.token_ids()
|
||||
}
|
||||
|
||||
/// Get a hash of the token IDs for caching purposes
|
||||
pub fn get_hash(&self) -> u64 {
|
||||
let mut hasher = DefaultHasher::new();
|
||||
self.hash(&mut hasher);
|
||||
hasher.finish()
|
||||
}
|
||||
}
|
||||
|
||||
/// Hash implementation for Encoding
|
||||
impl Hash for Encoding {
|
||||
fn hash<H: Hasher>(&self, state: &mut H) {
|
||||
match self {
|
||||
Encoding::Hf(inner) => inner.get_ids().hash(state),
|
||||
Encoding::Sp(inner) => inner.hash(state),
|
||||
Encoding::Tiktoken(inner) => inner.hash(state),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct SpecialTokens {
|
||||
pub bos_token: Option<String>,
|
||||
pub eos_token: Option<String>,
|
||||
pub unk_token: Option<String>,
|
||||
pub sep_token: Option<String>,
|
||||
pub pad_token: Option<String>,
|
||||
pub cls_token: Option<String>,
|
||||
pub mask_token: Option<String>,
|
||||
pub additional_special_tokens: Vec<String>,
|
||||
}
|
||||
@@ -1,312 +0,0 @@
|
||||
use smg::{
|
||||
protocols::chat::{ChatMessage, MessageContent},
|
||||
tokenizer::chat_template::{
|
||||
detect_chat_template_content_format, ChatTemplateContentFormat, ChatTemplateParams,
|
||||
ChatTemplateProcessor,
|
||||
},
|
||||
};
|
||||
|
||||
#[test]
|
||||
fn test_detect_string_format_deepseek() {
|
||||
// DeepSeek style template - expects string content
|
||||
let template = r#"
|
||||
{%- for message in messages %}
|
||||
{%- if message['role'] == 'user' %}
|
||||
User: {{ message['content'] }}
|
||||
{%- elif message['role'] == 'assistant' %}
|
||||
Assistant: {{ message['content'] }}
|
||||
{%- endif %}
|
||||
{%- endfor %}
|
||||
"#;
|
||||
|
||||
assert_eq!(
|
||||
detect_chat_template_content_format(template),
|
||||
ChatTemplateContentFormat::String
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_detect_openai_format_llama4() {
|
||||
// Llama4 style template - expects structured content
|
||||
let template = r#"
|
||||
{%- for message in messages %}
|
||||
{%- if message['content'] is iterable %}
|
||||
{%- for content in message['content'] %}
|
||||
{%- if content['type'] == 'text' %}
|
||||
{{ content['text'] }}
|
||||
{%- elif content['type'] == 'image' %}
|
||||
<image>
|
||||
{%- endif %}
|
||||
{%- endfor %}
|
||||
{%- else %}
|
||||
{{ message['content'] }}
|
||||
{%- endif %}
|
||||
{%- endfor %}
|
||||
"#;
|
||||
|
||||
assert_eq!(
|
||||
detect_chat_template_content_format(template),
|
||||
ChatTemplateContentFormat::OpenAI
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_detect_openai_format_dot_notation() {
|
||||
// Template using dot notation
|
||||
let template = r#"
|
||||
{%- for message in messages %}
|
||||
{%- for part in message.content %}
|
||||
{%- if part.type == 'text' %}
|
||||
{{ part.text }}
|
||||
{%- endif %}
|
||||
{%- endfor %}
|
||||
{%- endfor %}
|
||||
"#;
|
||||
|
||||
assert_eq!(
|
||||
detect_chat_template_content_format(template),
|
||||
ChatTemplateContentFormat::OpenAI
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_detect_openai_format_variable_assignment() {
|
||||
// Template that assigns content to variable then iterates
|
||||
let template = r#"
|
||||
{%- for message in messages %}
|
||||
{%- set content = message['content'] %}
|
||||
{%- if content is sequence %}
|
||||
{%- for item in content %}
|
||||
{{ item }}
|
||||
{%- endfor %}
|
||||
{%- endif %}
|
||||
{%- endfor %}
|
||||
"#;
|
||||
|
||||
assert_eq!(
|
||||
detect_chat_template_content_format(template),
|
||||
ChatTemplateContentFormat::OpenAI
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_detect_openai_format_glm4v_style() {
|
||||
// GLM4V uses 'msg' instead of 'message'
|
||||
let template = r#"
|
||||
{%- for msg in messages %}
|
||||
{%- for part in msg.content %}
|
||||
{%- if part.type == 'text' %}{{ part.text }}{%- endif %}
|
||||
{%- if part.type == 'image' %}<image>{%- endif %}
|
||||
{%- endfor %}
|
||||
{%- endfor %}
|
||||
"#;
|
||||
|
||||
assert_eq!(
|
||||
detect_chat_template_content_format(template),
|
||||
ChatTemplateContentFormat::OpenAI
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_detect_openai_format_with_length_check() {
|
||||
// Template that checks content length
|
||||
let template = r#"
|
||||
{%- for message in messages %}
|
||||
{%- if message.content|length > 0 %}
|
||||
{%- for item in message.content %}
|
||||
{{ item.text }}
|
||||
{%- endfor %}
|
||||
{%- endif %}
|
||||
{%- endfor %}
|
||||
"#;
|
||||
|
||||
assert_eq!(
|
||||
detect_chat_template_content_format(template),
|
||||
ChatTemplateContentFormat::OpenAI
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_detect_openai_format_with_index_access() {
|
||||
// Template that accesses content by index
|
||||
let template = r#"
|
||||
{%- for message in messages %}
|
||||
{%- if message.content[0] %}
|
||||
First item: {{ message.content[0].text }}
|
||||
{%- endif %}
|
||||
{%- endfor %}
|
||||
"#;
|
||||
|
||||
assert_eq!(
|
||||
detect_chat_template_content_format(template),
|
||||
ChatTemplateContentFormat::OpenAI
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_invalid_template_defaults_to_string() {
|
||||
let template = "Not a valid {% jinja template";
|
||||
|
||||
assert_eq!(
|
||||
detect_chat_template_content_format(template),
|
||||
ChatTemplateContentFormat::String
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_empty_template_defaults_to_string() {
|
||||
assert_eq!(
|
||||
detect_chat_template_content_format(""),
|
||||
ChatTemplateContentFormat::String
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_simple_chat_template_unit_test() {
|
||||
let template = r#"
|
||||
{%- for message in messages %}
|
||||
{{ message.role }}: {{ message.content }}
|
||||
{% endfor -%}
|
||||
{%- if add_generation_prompt %}
|
||||
assistant:
|
||||
{%- endif %}
|
||||
"#;
|
||||
|
||||
let processor = ChatTemplateProcessor::new(template.to_string());
|
||||
|
||||
let messages = [
|
||||
ChatMessage::System {
|
||||
content: MessageContent::Text("You are helpful".to_string()),
|
||||
name: None,
|
||||
},
|
||||
ChatMessage::User {
|
||||
content: MessageContent::Text("Hello".to_string()),
|
||||
name: None,
|
||||
},
|
||||
];
|
||||
|
||||
// Convert to JSON values like the router does
|
||||
let message_values: Vec<serde_json::Value> = messages
|
||||
.iter()
|
||||
.map(|msg| serde_json::to_value(msg).unwrap())
|
||||
.collect();
|
||||
|
||||
let params = ChatTemplateParams {
|
||||
add_generation_prompt: true,
|
||||
..Default::default()
|
||||
};
|
||||
let result = processor
|
||||
.apply_chat_template(&message_values, params)
|
||||
.unwrap();
|
||||
assert!(result.contains("system: You are helpful"));
|
||||
assert!(result.contains("user: Hello"));
|
||||
assert!(result.contains("assistant:"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_chat_template_with_tokens_unit_test() {
|
||||
// Template that uses template kwargs for tokens (more realistic)
|
||||
let template = r#"
|
||||
{%- if start_token -%}{{ start_token }}{%- endif -%}
|
||||
{%- for message in messages -%}
|
||||
{{ message.role }}: {{ message.content }}{%- if end_token -%}{{ end_token }}{%- endif -%}
|
||||
{% endfor -%}
|
||||
"#;
|
||||
|
||||
let processor = ChatTemplateProcessor::new(template.to_string());
|
||||
|
||||
let messages = [ChatMessage::User {
|
||||
content: MessageContent::Text("Test".to_string()),
|
||||
name: None,
|
||||
}];
|
||||
|
||||
// Convert to JSON values like the router does
|
||||
let message_values: Vec<serde_json::Value> = messages
|
||||
.iter()
|
||||
.map(|msg| serde_json::to_value(msg).unwrap())
|
||||
.collect();
|
||||
|
||||
// Use template_kwargs to pass tokens
|
||||
let mut template_kwargs = std::collections::HashMap::new();
|
||||
template_kwargs.insert(
|
||||
"start_token".to_string(),
|
||||
serde_json::Value::String("<s>".to_string()),
|
||||
);
|
||||
template_kwargs.insert(
|
||||
"end_token".to_string(),
|
||||
serde_json::Value::String("</s>".to_string()),
|
||||
);
|
||||
|
||||
let params = ChatTemplateParams {
|
||||
template_kwargs: Some(&template_kwargs),
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
let result = processor
|
||||
.apply_chat_template(&message_values, params)
|
||||
.unwrap();
|
||||
assert!(result.contains("<s>"));
|
||||
assert!(result.contains("</s>"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_detect_openai_format_qwen3vl_macro_style() {
|
||||
// Qwen3-VL style template using macros to handle multimodal content
|
||||
// This tests the macro-based detection pattern
|
||||
let template = r#"{%- set image_count = namespace(value=0) %}
|
||||
{%- set video_count = namespace(value=0) %}
|
||||
{%- macro render_content(content, do_vision_count) %}
|
||||
{%- if content is string %}
|
||||
{{- content }}
|
||||
{%- else %}
|
||||
{%- for item in content %}
|
||||
{%- if 'image' in item or 'image_url' in item or item.type == 'image' %}
|
||||
{%- if do_vision_count %}
|
||||
{%- set image_count.value = image_count.value + 1 %}
|
||||
{%- endif %}
|
||||
{%- if add_vision_id %}Picture {{ image_count.value }}: {% endif -%}
|
||||
<|vision_start|><|image_pad|><|vision_end|>
|
||||
{%- elif 'video' in item or item.type == 'video' %}
|
||||
{%- if do_vision_count %}
|
||||
{%- set video_count.value = video_count.value + 1 %}
|
||||
{%- endif %}
|
||||
{%- if add_vision_id %}Video {{ video_count.value }}: {% endif -%}
|
||||
<|vision_start|><|video_pad|><|vision_end|>
|
||||
{%- elif 'text' in item %}
|
||||
{{- item.text }}
|
||||
{%- endif %}
|
||||
{%- endfor %}
|
||||
{%- endif %}
|
||||
{%- endmacro %}
|
||||
{%- for message in messages %}
|
||||
{%- set content = render_content(message.content, True) %}
|
||||
{{- '<|im_start|>' + message.role + '\n' + content + '<|im_end|>' + '\n' }}
|
||||
{%- endfor %}
|
||||
{%- if add_generation_prompt %}
|
||||
{{- '<|im_start|>assistant\n' }}
|
||||
{%- endif %}"#;
|
||||
|
||||
assert_eq!(
|
||||
detect_chat_template_content_format(template),
|
||||
ChatTemplateContentFormat::OpenAI
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_detect_openai_format_arbitrary_variable_names() {
|
||||
// Test that detection works with any variable name, not just "message", "msg", "m"
|
||||
// Uses "chat_msg" and "x" as loop variables
|
||||
let template = r#"
|
||||
{%- for chat_msg in messages %}
|
||||
{%- for x in chat_msg.content %}
|
||||
{%- if x.type == 'text' %}{{ x.text }}{%- endif %}
|
||||
{%- if x.type == 'image' %}<image>{%- endif %}
|
||||
{%- endfor %}
|
||||
{%- endfor %}
|
||||
"#;
|
||||
|
||||
assert_eq!(
|
||||
detect_chat_template_content_format(template),
|
||||
ChatTemplateContentFormat::OpenAI
|
||||
);
|
||||
}
|
||||
@@ -1,414 +0,0 @@
|
||||
use smg::{
|
||||
protocols::{
|
||||
chat::{ChatMessage, MessageContent},
|
||||
common::{ContentPart, ImageUrl},
|
||||
},
|
||||
tokenizer::chat_template::{
|
||||
detect_chat_template_content_format, ChatTemplateContentFormat, ChatTemplateParams,
|
||||
ChatTemplateProcessor,
|
||||
},
|
||||
};
|
||||
|
||||
#[test]
|
||||
fn test_simple_chat_template() {
|
||||
let template = r#"
|
||||
{%- for message in messages %}
|
||||
<|{{ message.role }}|>{{ message.content }}<|end|>
|
||||
{% endfor -%}
|
||||
{%- if add_generation_prompt %}
|
||||
<|assistant|>
|
||||
{%- endif %}
|
||||
"#;
|
||||
|
||||
let processor = ChatTemplateProcessor::new(template.to_string());
|
||||
|
||||
let messages = [ChatMessage::User {
|
||||
content: MessageContent::Text("Test".to_string()),
|
||||
name: None,
|
||||
}];
|
||||
|
||||
// Convert to JSON values like the router does
|
||||
let message_values: Vec<serde_json::Value> = messages
|
||||
.iter()
|
||||
.map(|msg| serde_json::to_value(msg).unwrap())
|
||||
.collect();
|
||||
|
||||
let params = ChatTemplateParams {
|
||||
add_generation_prompt: true,
|
||||
..Default::default()
|
||||
};
|
||||
let result = processor
|
||||
.apply_chat_template(&message_values, params)
|
||||
.unwrap();
|
||||
assert!(result.contains("<|user|>Test<|end|>"));
|
||||
assert!(result.contains("<|assistant|>"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_chat_template_with_tokens() {
|
||||
// Template that uses template kwargs for tokens
|
||||
let template = r#"
|
||||
{%- if bos_token -%}{{ bos_token }}{%- endif -%}
|
||||
{%- for message in messages -%}
|
||||
{{ message.role }}: {{ message.content }}{%- if eos_token -%}{{ eos_token }}{%- endif -%}
|
||||
{% endfor -%}
|
||||
"#;
|
||||
|
||||
let processor = ChatTemplateProcessor::new(template.to_string());
|
||||
|
||||
let messages = [ChatMessage::User {
|
||||
content: MessageContent::Text("Test".to_string()),
|
||||
name: None,
|
||||
}];
|
||||
|
||||
// Convert to JSON values like the router does
|
||||
let message_values: Vec<serde_json::Value> = messages
|
||||
.iter()
|
||||
.map(|msg| serde_json::to_value(msg).unwrap())
|
||||
.collect();
|
||||
|
||||
// Use template_kwargs to pass tokens
|
||||
let mut template_kwargs = std::collections::HashMap::new();
|
||||
template_kwargs.insert(
|
||||
"bos_token".to_string(),
|
||||
serde_json::Value::String("<s>".to_string()),
|
||||
);
|
||||
template_kwargs.insert(
|
||||
"eos_token".to_string(),
|
||||
serde_json::Value::String("</s>".to_string()),
|
||||
);
|
||||
|
||||
let params = ChatTemplateParams {
|
||||
template_kwargs: Some(&template_kwargs),
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
let result = processor
|
||||
.apply_chat_template(&message_values, params)
|
||||
.unwrap();
|
||||
assert!(result.contains("<s>"));
|
||||
assert!(result.contains("</s>"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_llama_style_template() {
|
||||
let template = r#"
|
||||
{%- if messages[0]['role'] == 'system' -%}
|
||||
{%- set system_message = messages[0]['content'] -%}
|
||||
{%- set messages = messages[1:] -%}
|
||||
{%- else -%}
|
||||
{%- set system_message = '' -%}
|
||||
{%- endif -%}
|
||||
|
||||
{{- bos_token if bos_token else '<|begin_of_text|>' }}
|
||||
{%- if system_message %}
|
||||
{{- '<|start_header_id|>system<|end_header_id|>\n\n' + system_message + '<|eot_id|>' }}
|
||||
{%- endif %}
|
||||
|
||||
{%- for message in messages %}
|
||||
{{- '<|start_header_id|>' + message['role'] + '<|end_header_id|>\n\n' + message['content'] + '<|eot_id|>' }}
|
||||
{%- endfor %}
|
||||
|
||||
{%- if add_generation_prompt %}
|
||||
{{- '<|start_header_id|>assistant<|end_header_id|>\n\n' }}
|
||||
{%- endif %}
|
||||
"#;
|
||||
|
||||
let processor = ChatTemplateProcessor::new(template.to_string());
|
||||
|
||||
let messages = [
|
||||
ChatMessage::System {
|
||||
content: MessageContent::Text("You are a helpful assistant".to_string()),
|
||||
name: None,
|
||||
},
|
||||
ChatMessage::User {
|
||||
content: MessageContent::Text("What is 2+2?".to_string()),
|
||||
name: None,
|
||||
},
|
||||
];
|
||||
|
||||
// Convert to JSON values
|
||||
let json_messages: Vec<serde_json::Value> = messages
|
||||
.iter()
|
||||
.map(|msg| serde_json::to_value(msg).unwrap())
|
||||
.collect();
|
||||
|
||||
// Use template_kwargs to pass the token
|
||||
let mut template_kwargs = std::collections::HashMap::new();
|
||||
template_kwargs.insert(
|
||||
"bos_token".to_string(),
|
||||
serde_json::Value::String("<|begin_of_text|>".to_string()),
|
||||
);
|
||||
|
||||
let params = ChatTemplateParams {
|
||||
add_generation_prompt: true,
|
||||
template_kwargs: Some(&template_kwargs),
|
||||
..Default::default()
|
||||
};
|
||||
let result = processor
|
||||
.apply_chat_template(&json_messages, params)
|
||||
.unwrap();
|
||||
|
||||
// Check that the result contains expected markers
|
||||
assert!(result.contains("<|begin_of_text|>"));
|
||||
assert!(result.contains("<|start_header_id|>system<|end_header_id|>"));
|
||||
assert!(result.contains("You are a helpful assistant"));
|
||||
assert!(result.contains("<|start_header_id|>user<|end_header_id|>"));
|
||||
assert!(result.contains("What is 2+2?"));
|
||||
assert!(result.contains("<|start_header_id|>assistant<|end_header_id|>"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_chatml_template() {
|
||||
let template = r#"
|
||||
{%- for message in messages %}
|
||||
{{- '<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>\n' }}
|
||||
{%- endfor %}
|
||||
{%- if add_generation_prompt %}
|
||||
{{- '<|im_start|>assistant\n' }}
|
||||
{%- endif %}
|
||||
"#;
|
||||
|
||||
let processor = ChatTemplateProcessor::new(template.to_string());
|
||||
|
||||
let messages = [
|
||||
ChatMessage::User {
|
||||
content: MessageContent::Text("Hello".to_string()),
|
||||
name: None,
|
||||
},
|
||||
ChatMessage::Assistant {
|
||||
content: Some(MessageContent::Text("Hi there!".to_string())),
|
||||
name: None,
|
||||
tool_calls: None,
|
||||
reasoning_content: None,
|
||||
},
|
||||
ChatMessage::User {
|
||||
content: MessageContent::Text("How are you?".to_string()),
|
||||
name: None,
|
||||
},
|
||||
];
|
||||
|
||||
// Convert to JSON values
|
||||
let json_messages: Vec<serde_json::Value> = messages
|
||||
.iter()
|
||||
.map(|msg| serde_json::to_value(msg).unwrap())
|
||||
.collect();
|
||||
|
||||
let result = processor
|
||||
.apply_chat_template(
|
||||
&json_messages,
|
||||
ChatTemplateParams {
|
||||
add_generation_prompt: true,
|
||||
..Default::default()
|
||||
},
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
// Check ChatML format
|
||||
assert!(result.contains("<|im_start|>user\nHello<|im_end|>"));
|
||||
assert!(result.contains("<|im_start|>assistant\nHi there!<|im_end|>"));
|
||||
assert!(result.contains("<|im_start|>user\nHow are you?<|im_end|>"));
|
||||
assert!(result.ends_with("<|im_start|>assistant\n"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_template_without_generation_prompt() {
|
||||
let template = r#"
|
||||
{%- for message in messages -%}
|
||||
{{ message.role }}: {{ message.content }}
|
||||
{% endfor -%}
|
||||
{%- if add_generation_prompt -%}
|
||||
assistant:
|
||||
{%- endif -%}
|
||||
"#;
|
||||
|
||||
let processor = ChatTemplateProcessor::new(template.to_string());
|
||||
|
||||
let messages = [ChatMessage::User {
|
||||
content: MessageContent::Text("Test".to_string()),
|
||||
name: None,
|
||||
}];
|
||||
|
||||
// Convert to JSON values
|
||||
let json_messages: Vec<serde_json::Value> = messages
|
||||
.iter()
|
||||
.map(|msg| serde_json::to_value(msg).unwrap())
|
||||
.collect();
|
||||
|
||||
let result = processor
|
||||
.apply_chat_template(&json_messages, ChatTemplateParams::default())
|
||||
.unwrap();
|
||||
assert_eq!(result.trim(), "user: Test");
|
||||
|
||||
let result_with_prompt = processor
|
||||
.apply_chat_template(
|
||||
&json_messages,
|
||||
ChatTemplateParams {
|
||||
add_generation_prompt: true,
|
||||
..Default::default()
|
||||
},
|
||||
)
|
||||
.unwrap();
|
||||
assert!(result_with_prompt.contains("assistant:"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_empty_messages_template() {
|
||||
let template = r#"{% for msg in messages %}{{ msg.role }}: {{ msg.content }}\n{% endfor %}"#;
|
||||
|
||||
let processor = ChatTemplateProcessor::new(template.to_string());
|
||||
|
||||
let messages: Vec<serde_json::Value> = vec![];
|
||||
let result = processor
|
||||
.apply_chat_template(&messages, ChatTemplateParams::default())
|
||||
.unwrap();
|
||||
assert_eq!(result, "");
|
||||
}
|
||||
|
||||
/// Test that tojson filter accepts ensure_ascii kwarg (HuggingFace compatibility)
|
||||
/// This is the fix for: "unknown keyword argument 'ensure_ascii'"
|
||||
#[test]
|
||||
fn test_tojson_with_ensure_ascii() {
|
||||
// Template that uses tojson(ensure_ascii=False) like HuggingFace templates do
|
||||
let template = r#"
|
||||
{%- for message in messages -%}
|
||||
{{ message.role }}: {{ message.content }}
|
||||
{%- if message.tool_calls is defined and message.tool_calls -%}
|
||||
Tools: {{ message.tool_calls|tojson(ensure_ascii=False) }}
|
||||
{%- endif -%}
|
||||
{% endfor -%}
|
||||
"#;
|
||||
|
||||
let processor = ChatTemplateProcessor::new(template.to_string());
|
||||
|
||||
let messages = [ChatMessage::User {
|
||||
content: MessageContent::Text("Test with Unicode: 日本語".to_string()),
|
||||
name: None,
|
||||
}];
|
||||
|
||||
// Convert to JSON values
|
||||
let json_messages: Vec<serde_json::Value> = messages
|
||||
.iter()
|
||||
.map(|msg| serde_json::to_value(msg).unwrap())
|
||||
.collect();
|
||||
|
||||
// This should NOT fail with "unknown keyword argument 'ensure_ascii'"
|
||||
let result = processor
|
||||
.apply_chat_template(&json_messages, ChatTemplateParams::default())
|
||||
.unwrap();
|
||||
|
||||
assert!(result.contains("user: Test with Unicode: 日本語"));
|
||||
}
|
||||
|
||||
/// Test tojson with all HuggingFace kwargs
|
||||
#[test]
|
||||
fn test_tojson_with_all_huggingface_kwargs() {
|
||||
// Template using all the kwargs that HuggingFace's custom tojson accepts
|
||||
let template = r#"
|
||||
{%- set data = {"z_key": 1, "a_key": 2, "m_key": 3} -%}
|
||||
Unsorted: {{ data|tojson }}
|
||||
Sorted: {{ data|tojson(sort_keys=True) }}
|
||||
Indented: {{ data|tojson(indent=2) }}
|
||||
All: {{ data|tojson(ensure_ascii=False, sort_keys=True, indent=2) }}
|
||||
"#;
|
||||
|
||||
let processor = ChatTemplateProcessor::new(template.to_string());
|
||||
let messages: Vec<serde_json::Value> = vec![];
|
||||
|
||||
// This should NOT fail - all kwargs should be accepted
|
||||
let result = processor
|
||||
.apply_chat_template(&messages, ChatTemplateParams::default())
|
||||
.unwrap();
|
||||
|
||||
// Verify sorted output contains keys in alphabetical order
|
||||
assert!(result.contains("Sorted:"));
|
||||
// The sorted output should have a_key before m_key before z_key
|
||||
let sorted_line = result.lines().find(|l| l.starts_with("Sorted:")).unwrap();
|
||||
let a_pos = sorted_line.find("a_key").unwrap();
|
||||
let m_pos = sorted_line.find("m_key").unwrap();
|
||||
let z_pos = sorted_line.find("z_key").unwrap();
|
||||
assert!(a_pos < m_pos && m_pos < z_pos, "Keys should be sorted");
|
||||
|
||||
// Verify indented output is pretty-printed with newlines
|
||||
assert!(
|
||||
result.contains("Indented: {\n"),
|
||||
"Indented JSON should be pretty-printed with newlines"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_content_format_detection() {
|
||||
let string_template = r#"
|
||||
{%- for message in messages -%}
|
||||
{{ message.role }}: {{ message.content }}
|
||||
{%- endfor -%}
|
||||
"#;
|
||||
assert_eq!(
|
||||
detect_chat_template_content_format(string_template),
|
||||
ChatTemplateContentFormat::String
|
||||
);
|
||||
|
||||
let openai_template = r#"
|
||||
{%- for message in messages -%}
|
||||
{%- for content in message.content -%}
|
||||
{{ content.type }}: {{ content.text }}
|
||||
{%- endfor -%}
|
||||
{%- endfor -%}
|
||||
"#;
|
||||
assert_eq!(
|
||||
detect_chat_template_content_format(openai_template),
|
||||
ChatTemplateContentFormat::OpenAI
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_template_with_multimodal_content() {
|
||||
let template = r#"
|
||||
{%- for message in messages %}
|
||||
{{ message.role }}:
|
||||
{%- if message.content is string %}
|
||||
{{ message.content }}
|
||||
{%- else %}
|
||||
{%- for part in message.content %}
|
||||
{%- if part.type == "text" %}
|
||||
{{ part.text }}
|
||||
{%- elif part.type == "image_url" %}
|
||||
[IMAGE]
|
||||
{%- endif %}
|
||||
{%- endfor %}
|
||||
{%- endif %}
|
||||
{% endfor %}
|
||||
"#;
|
||||
|
||||
let processor = ChatTemplateProcessor::new(template.to_string());
|
||||
|
||||
let messages = [ChatMessage::User {
|
||||
content: MessageContent::Parts(vec![
|
||||
ContentPart::Text {
|
||||
text: "Look at this:".to_string(),
|
||||
},
|
||||
ContentPart::ImageUrl {
|
||||
image_url: ImageUrl {
|
||||
url: "https://example.com/image.jpg".to_string(),
|
||||
detail: None,
|
||||
},
|
||||
},
|
||||
]),
|
||||
name: None,
|
||||
}];
|
||||
|
||||
// Convert to JSON values
|
||||
let json_messages: Vec<serde_json::Value> = messages
|
||||
.iter()
|
||||
.map(|msg| serde_json::to_value(msg).unwrap())
|
||||
.collect();
|
||||
|
||||
let result = processor
|
||||
.apply_chat_template(&json_messages, ChatTemplateParams::default())
|
||||
.unwrap();
|
||||
|
||||
// Should contain both text and image parts
|
||||
assert!(result.contains("user:"));
|
||||
assert!(result.contains("Look at this:"));
|
||||
assert!(result.contains("[IMAGE]"));
|
||||
}
|
||||
@@ -1,230 +0,0 @@
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::fs;
|
||||
|
||||
use smg::{
|
||||
protocols::chat::{ChatMessage, MessageContent},
|
||||
tokenizer::{chat_template::ChatTemplateParams, huggingface::HuggingFaceTokenizer},
|
||||
};
|
||||
use tempfile::TempDir;
|
||||
|
||||
#[test]
|
||||
fn test_load_chat_template_from_file() {
|
||||
// Create temporary directory
|
||||
let temp_dir = TempDir::new().unwrap();
|
||||
let template_path = temp_dir.path().join("template.jinja");
|
||||
|
||||
// Write a test template
|
||||
let template_content = r#"
|
||||
{%- for message in messages %}
|
||||
{{- '<|' + message['role'] + '|>' + message['content'] }}
|
||||
{%- endfor %}
|
||||
{%- if add_generation_prompt %}
|
||||
{{- '<|assistant|>' }}
|
||||
{%- endif %}
|
||||
"#;
|
||||
fs::write(&template_path, template_content).unwrap();
|
||||
|
||||
// Create a mock tokenizer config
|
||||
let tokenizer_config = r#"{
|
||||
"version": "1.0",
|
||||
"truncation": null,
|
||||
"padding": null,
|
||||
"added_tokens": [],
|
||||
"normalizer": null,
|
||||
"pre_tokenizer": {
|
||||
"type": "Whitespace"
|
||||
},
|
||||
"post_processor": null,
|
||||
"decoder": null,
|
||||
"model": {
|
||||
"type": "BPE",
|
||||
"vocab": {
|
||||
"hello": 0,
|
||||
"world": 1,
|
||||
"<s>": 2,
|
||||
"</s>": 3
|
||||
},
|
||||
"merges": []
|
||||
}
|
||||
}"#;
|
||||
|
||||
let tokenizer_path = temp_dir.path().join("tokenizer.json");
|
||||
fs::write(&tokenizer_path, tokenizer_config).unwrap();
|
||||
|
||||
// Load tokenizer with custom chat template
|
||||
let tokenizer = HuggingFaceTokenizer::from_file_with_chat_template(
|
||||
tokenizer_path.to_str().unwrap(),
|
||||
Some(template_path.to_str().unwrap()),
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let messages = [
|
||||
ChatMessage::User {
|
||||
content: MessageContent::Text("Hello".to_string()),
|
||||
name: None,
|
||||
},
|
||||
ChatMessage::Assistant {
|
||||
content: Some(MessageContent::Text("Hi there".to_string())),
|
||||
name: None,
|
||||
tool_calls: None,
|
||||
reasoning_content: None,
|
||||
},
|
||||
];
|
||||
|
||||
// Convert to JSON values like the router does
|
||||
let json_messages: Vec<serde_json::Value> = messages
|
||||
.iter()
|
||||
.map(|msg| serde_json::to_value(msg).unwrap())
|
||||
.collect();
|
||||
|
||||
use smg::tokenizer::chat_template::ChatTemplateParams;
|
||||
let params = ChatTemplateParams {
|
||||
add_generation_prompt: true,
|
||||
..Default::default()
|
||||
};
|
||||
let result = tokenizer
|
||||
.apply_chat_template(&json_messages, params)
|
||||
.unwrap();
|
||||
|
||||
assert!(result.contains("<|user|>Hello"));
|
||||
assert!(result.contains("<|assistant|>Hi there"));
|
||||
assert!(result.ends_with("<|assistant|>"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_override_existing_template() {
|
||||
// Create temporary directory
|
||||
let temp_dir = TempDir::new().unwrap();
|
||||
|
||||
// Create tokenizer config with a built-in template
|
||||
let tokenizer_config_path = temp_dir.path().join("tokenizer_config.json");
|
||||
let config_with_template = r#"{
|
||||
"chat_template": "built-in: {% for msg in messages %}{{ msg.content }}{% endfor %}"
|
||||
}"#;
|
||||
fs::write(&tokenizer_config_path, config_with_template).unwrap();
|
||||
|
||||
// Create the actual tokenizer file
|
||||
let tokenizer_json = r#"{
|
||||
"version": "1.0",
|
||||
"truncation": null,
|
||||
"padding": null,
|
||||
"added_tokens": [],
|
||||
"normalizer": null,
|
||||
"pre_tokenizer": {
|
||||
"type": "Whitespace"
|
||||
},
|
||||
"post_processor": null,
|
||||
"decoder": null,
|
||||
"model": {
|
||||
"type": "BPE",
|
||||
"vocab": {
|
||||
"test": 0,
|
||||
"<s>": 1,
|
||||
"</s>": 2
|
||||
},
|
||||
"merges": []
|
||||
}
|
||||
}"#;
|
||||
let tokenizer_path = temp_dir.path().join("tokenizer.json");
|
||||
fs::write(&tokenizer_path, tokenizer_json).unwrap();
|
||||
|
||||
// Create custom template that should override
|
||||
let custom_template_path = temp_dir.path().join("custom.jinja");
|
||||
let custom_template =
|
||||
r#"CUSTOM: {% for msg in messages %}[{{ msg.role }}]: {{ msg.content }}{% endfor %}"#;
|
||||
fs::write(&custom_template_path, custom_template).unwrap();
|
||||
|
||||
// Load with custom template - should override the built-in one
|
||||
let tokenizer = HuggingFaceTokenizer::from_file_with_chat_template(
|
||||
tokenizer_path.to_str().unwrap(),
|
||||
Some(custom_template_path.to_str().unwrap()),
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let messages = [ChatMessage::User {
|
||||
content: MessageContent::Text("Test".to_string()),
|
||||
name: None,
|
||||
}];
|
||||
|
||||
// Convert to JSON values
|
||||
let json_messages: Vec<serde_json::Value> = messages
|
||||
.iter()
|
||||
.map(|msg| serde_json::to_value(msg).unwrap())
|
||||
.collect();
|
||||
|
||||
let result = tokenizer
|
||||
.apply_chat_template(&json_messages, ChatTemplateParams::default())
|
||||
.unwrap();
|
||||
|
||||
// Should use CUSTOM template, not built-in
|
||||
assert!(result.starts_with("CUSTOM:"));
|
||||
assert!(result.contains("[user]: Test"));
|
||||
assert!(!result.contains("built-in:"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_set_chat_template_after_creation() {
|
||||
// Create temporary directory and tokenizer file
|
||||
let temp_dir = TempDir::new().unwrap();
|
||||
let tokenizer_json = r#"{
|
||||
"version": "1.0",
|
||||
"truncation": null,
|
||||
"padding": null,
|
||||
"added_tokens": [],
|
||||
"normalizer": null,
|
||||
"pre_tokenizer": {
|
||||
"type": "Whitespace"
|
||||
},
|
||||
"post_processor": null,
|
||||
"decoder": null,
|
||||
"model": {
|
||||
"type": "BPE",
|
||||
"vocab": {
|
||||
"test": 0,
|
||||
"<s>": 1,
|
||||
"</s>": 2
|
||||
},
|
||||
"merges": []
|
||||
}
|
||||
}"#;
|
||||
let tokenizer_path = temp_dir.path().join("tokenizer.json");
|
||||
fs::write(&tokenizer_path, tokenizer_json).unwrap();
|
||||
|
||||
// Load tokenizer without custom template
|
||||
let mut tokenizer =
|
||||
HuggingFaceTokenizer::from_file(tokenizer_path.to_str().unwrap()).unwrap();
|
||||
|
||||
// Set a template after creation (mimics Python's behavior)
|
||||
let new_template =
|
||||
"NEW: {% for msg in messages %}{{ msg.role }}: {{ msg.content }}; {% endfor %}";
|
||||
tokenizer.set_chat_template(new_template.to_string());
|
||||
|
||||
let messages = [
|
||||
ChatMessage::User {
|
||||
content: MessageContent::Text("Hello".to_string()),
|
||||
name: None,
|
||||
},
|
||||
ChatMessage::Assistant {
|
||||
content: Some(MessageContent::Text("World".to_string())),
|
||||
name: None,
|
||||
tool_calls: None,
|
||||
reasoning_content: None,
|
||||
},
|
||||
];
|
||||
|
||||
// Convert to JSON values
|
||||
let json_messages: Vec<serde_json::Value> = messages
|
||||
.iter()
|
||||
.map(|msg| serde_json::to_value(msg).unwrap())
|
||||
.collect();
|
||||
|
||||
let result = tokenizer
|
||||
.apply_chat_template(&json_messages, ChatTemplateParams::default())
|
||||
.unwrap();
|
||||
|
||||
assert!(result.starts_with("NEW:"));
|
||||
assert!(result.contains("user: Hello;"));
|
||||
assert!(result.contains("assistant: World;"));
|
||||
}
|
||||
}
|
||||
@@ -1,7 +0,0 @@
|
||||
//! Tokenizer and chat template integration tests
|
||||
|
||||
mod chat_template_format_detection;
|
||||
mod chat_template_integration;
|
||||
mod chat_template_loading;
|
||||
mod tokenizer_cache_correctness_test;
|
||||
mod tokenizer_integration;
|
||||
@@ -1,471 +0,0 @@
|
||||
//! Cache correctness integration test
|
||||
//!
|
||||
//! This test validates that the tokenizer cache (L0, L1, and L0+L1 combined) produces
|
||||
//! exactly the same token IDs as uncached tokenization across multiple chat turns.
|
||||
//! Uses the real Qwen/Qwen3-4B-Instruct-2507 tokenizer to test with actual special tokens.
|
||||
|
||||
use std::{
|
||||
path::PathBuf,
|
||||
sync::{Arc, OnceLock},
|
||||
};
|
||||
|
||||
use smg::tokenizer::{
|
||||
cache::{CacheConfig, CachedTokenizer},
|
||||
hub::download_tokenizer_from_hf,
|
||||
huggingface::HuggingFaceTokenizer,
|
||||
traits::Encoder,
|
||||
};
|
||||
|
||||
/// Global tokenizer path cache - download once, reuse across all tests
|
||||
static TOKENIZER_PATH: OnceLock<Option<PathBuf>> = OnceLock::new();
|
||||
|
||||
/// Download Qwen3-4B-Instruct-2507 tokenizer once and cache the path
|
||||
async fn get_tokenizer_path() -> Option<PathBuf> {
|
||||
// Check if already downloaded
|
||||
if let Some(cached) = TOKENIZER_PATH.get() {
|
||||
return cached.clone();
|
||||
}
|
||||
|
||||
// Download tokenizer
|
||||
let result = match download_tokenizer_from_hf("Qwen/Qwen3-4B-Instruct-2507").await {
|
||||
Ok(cache_dir) => {
|
||||
let tokenizer_path = cache_dir.join("tokenizer.json");
|
||||
if tokenizer_path.exists() {
|
||||
Some(tokenizer_path)
|
||||
} else {
|
||||
println!("Tokenizer downloaded but tokenizer.json not found");
|
||||
None
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
println!("Failed to download tokenizer: {}", e);
|
||||
None
|
||||
}
|
||||
};
|
||||
|
||||
// Cache the result (even if None, so we don't retry on failure)
|
||||
TOKENIZER_PATH.set(result.clone()).ok();
|
||||
result
|
||||
}
|
||||
|
||||
/// Comprehensive multi-turn chat conversation for testing cache correctness
|
||||
/// Uses Qwen's special tokens with diverse content to hit edge cases
|
||||
const CHAT_TURNS: [&str; 29] = [
|
||||
// Basic conversation patterns
|
||||
"<|im_start|>system\nYou are a helpful AI assistant.<|im_end|>",
|
||||
"<|im_start|>system\nYou are a helpful AI assistant.<|im_end|><|im_start|>user\nWhat is the capital of France?<|im_end|>",
|
||||
"<|im_start|>system\nYou are a helpful AI assistant.<|im_end|><|im_start|>user\nWhat is the capital of France?<|im_end|><|im_start|>assistant\nThe capital of France is Paris.<|im_end|>",
|
||||
|
||||
// Different system prompts (testing different prefix patterns)
|
||||
"<|im_start|>system\nYou are a coding tutor specializing in Rust programming.<|im_end|><|im_start|>user\nExplain ownership.<|im_end|>",
|
||||
"<|im_start|>system\nYou are a math teacher.<|im_end|><|im_start|>user\nSolve: 2x + 5 = 13<|im_end|>",
|
||||
|
||||
// Long conversation with multiple turns (testing longer prefixes)
|
||||
"<|im_start|>system\nYou are a helpful AI assistant.<|im_end|><|im_start|>user\nTell me about deep learning.<|im_end|><|im_start|>assistant\nDeep learning is a subset of machine learning that uses neural networks with multiple layers.<|im_end|><|im_start|>user\nWhat are the main architectures?<|im_end|>",
|
||||
|
||||
// Code snippets (testing different character patterns)
|
||||
"<|im_start|>system\nYou are a code reviewer.<|im_end|><|im_start|>user\nReview this code:\nfn main() {\n println!(\"Hello, world!\");\n}\n<|im_end|>",
|
||||
"<|im_start|>system\nYou are a code reviewer.<|im_end|><|im_start|>user\nExplain this Rust code:\nimpl<T> Drop for Box<T> {\n fn drop(&mut self) { /* ... */ }\n}\n<|im_end|>",
|
||||
|
||||
// Mathematical content
|
||||
"<|im_start|>system\nYou are a math tutor.<|im_end|><|im_start|>user\nProve that √2 is irrational using proof by contradiction.<|im_end|>",
|
||||
"<|im_start|>system\nYou are a math tutor.<|im_end|><|im_start|>user\nCalculate: ∫(x² + 3x + 2)dx from 0 to 5<|im_end|>",
|
||||
|
||||
// Multilingual content
|
||||
"<|im_start|>system\nYou are a multilingual assistant.<|im_end|><|im_start|>user\nTranslate to French: The quick brown fox jumps over the lazy dog.<|im_end|>",
|
||||
"<|im_start|>system\nYou are a multilingual assistant.<|im_end|><|im_start|>user\n你好，请帮我翻译这句话：I love programming in Rust.<|im_end|>",
|
||||
"<|im_start|>system\nYou are a multilingual assistant.<|im_end|><|im_start|>user\nこんにちは！Rustについて教えてください。<|im_end|>",
|
||||
|
||||
// Special characters and emojis
|
||||
"<|im_start|>system\nYou are a friendly chatbot.<|im_end|><|im_start|>user\nWhat do you think about emojis? 😀🎉🚀💻<|im_end|>",
|
||||
"<|im_start|>system\nYou are a data analyst.<|im_end|><|im_start|>user\nAnalyze this: {\"name\": \"test\", \"value\": 42, \"nested\": {\"key\": \"value\"}}<|im_end|>",
|
||||
|
||||
// Very long message (testing large token counts)
|
||||
"<|im_start|>system\nYou are a literature expert.<|im_end|><|im_start|>user\nAnalyze the themes in this passage: In the vast expanse of the digital realm, where bits and bytes dance in harmonious symphony, there exists a paradigm that transcends mere computation. This paradigm, known as machine learning, represents humanity's quest to imbue silicon with the spark of cognition. Deep neural networks, inspired by the intricate architecture of biological brains, layer upon layer of artificial neurons, each connection a synapse firing in the dark recesses of mathematical space. Through gradient descent, these networks learn patterns invisible to human perception, extracting meaning from chaos, signal from noise. The transformer architecture revolutionized this field, introducing attention mechanisms that allowed models to focus on relevant information, much like how humans selectively attend to important details in their environment.<|im_end|>",
|
||||
|
||||
// Edge case: Multiple special tokens in sequence
|
||||
"<|im_start|>system\nYou are helpful.<|im_end|><|im_start|>user\nHi<|im_end|><|im_start|>assistant\nHello!<|im_end|><|im_start|>user\nHow are you?<|im_end|>",
|
||||
|
||||
// Edge case: Empty-ish messages
|
||||
"<|im_start|>system\n<|im_end|><|im_start|>user\nTest<|im_end|>",
|
||||
"<|im_start|>system\nBrief.<|im_end|><|im_start|>user\nOK<|im_end|>",
|
||||
|
||||
// Technical documentation style
|
||||
"<|im_start|>system\nYou are a technical writer.<|im_end|><|im_start|>user\nDocument the following API:\n\n```rust\npub struct CachedTokenizer {\n inner: Arc<dyn Tokenizer>,\n l0: Option<L0Cache>,\n l1: Option<L1Cache>,\n}\n\nimpl Encoder for CachedTokenizer {\n fn encode(&self, input: &str) -> Result<Encoding>;\n}\n```\n<|im_end|>",
|
||||
|
||||
// Conversation with code review
|
||||
"<|im_start|>system\nYou are a senior Rust developer.<|im_end|><|im_start|>user\nReview for correctness:\n\nlet special_tokens: Option<Vec<&str>> = self.l1.as_ref().map(|_| {\n self.special_token_strings.iter().map(|s| s.as_str()).collect()\n});<|im_end|>",
|
||||
|
||||
// Markdown formatted content
|
||||
"<|im_start|>system\nYou are a documentation assistant.<|im_end|><|im_start|>user\nFormat this as markdown:\n\n# Cache Architecture\n\n## L0 Cache\n- Exact match\n- DashMap based\n- 10K entries\n\n## L1 Cache \n- Prefix match\n- Special token boundaries\n- 50MB memory\n<|im_end|>",
|
||||
|
||||
// Complex nested structures
|
||||
"<|im_start|>system\nYou are a JSON expert.<|im_end|><|im_start|>user\nValidate this JSON:\n{\n \"tokenizer_cache\": {\n \"enable_l0\": true,\n \"l0_max_entries\": 10000,\n \"enable_l1\": true,\n \"l1_max_memory\": 52428800,\n \"stats\": {\n \"hits\": [1, 2, 3],\n \"misses\": {\"count\": 5}\n }\n }\n}\n<|im_end|>",
|
||||
|
||||
// SQL queries
|
||||
"<|im_start|>system\nYou are a database expert.<|im_end|><|im_start|>user\nOptimize this query:\nSELECT u.name, COUNT(p.id) as post_count\nFROM users u\nLEFT JOIN posts p ON u.id = p.user_id\nWHERE u.created_at > '2024-01-01'\nGROUP BY u.id, u.name\nHAVING COUNT(p.id) > 5\nORDER BY post_count DESC;<|im_end|>",
|
||||
|
||||
// Regex patterns
|
||||
"<|im_start|>system\nYou are a regex expert.<|im_end|><|im_start|>user\nExplain this regex: ^(?:[a-zA-Z0-9](?:[a-zA-Z0-9-]{0,61}[a-zA-Z0-9])?\\.)+[a-zA-Z]{2,}$<|im_end|>",
|
||||
|
||||
// Command line examples
|
||||
"<|im_start|>system\nYou are a DevOps engineer.<|im_end|><|im_start|>user\nExplain this command:\ncargo bench --bench tokenizer_benchmark -- --color=never | tee results.txt<|im_end|>",
|
||||
|
||||
// Unicode edge cases
|
||||
"<|im_start|>system\nYou are helpful.<|im_end|><|im_start|>user\nTest: café, naïve, Zürich, 北京, 東京, मुंबई, Москва<|im_end|>",
|
||||
|
||||
// Mixed content complexity
|
||||
"<|im_start|>system\nYou are a software architect.<|im_end|><|im_start|>user\nDesign a caching system that:\n1. Handles 10K+ QPS\n2. Maintains 99.9% uptime \n3. Supports L0 (exact) and L1 (prefix) caching\n4. Uses Blake3 for hashing (10GB/s throughput)\n5. Implements LRU eviction\n6. Thread-safe with lock-free reads\n\nKey requirements:\n- Memory: 50MB L1 budget\n- Latency: <100µs p99\n- Correctness: 100% (no false tokens)\n<|im_end|>",
|
||||
|
||||
// Very long technical discussion
|
||||
"<|im_start|>system\nYou are a compiler expert.<|im_end|><|im_start|>user\nExplain why BPE tokenizers are not prefix-stable:\n\nThe core issue is that BPE applies merges based on local context. When you tokenize 'prefix' alone, it might apply merge rules differently than when tokenizing 'prefix + suffix' as a whole. For example:\n\ntokenize('hello world') might produce [hello, _world]\ntokenize('hello') + tokenize(' world') might produce [hel, lo, _wo, rld]\n\nThis is because the merge rules see different contexts. The space before 'world' in the first case is part of the token boundary, but in the second case, ' world' is tokenized in isolation.\n\nSpecial tokens solve this because they are:\n1. Atomic (never split or merged)\n2. Protected from normalization\n3. Marked with special: true flag\n4. Have normalized: false property\n\nThis guarantees: tokenize(prefix + special + suffix) = tokenize(prefix + special) + tokenize(suffix)\n\nOur L1 cache exploits this by:\n1. Finding all special token boundaries\n2. Re-tokenizing prefixes at those boundaries\n3. Caching the exact token IDs\n4. On cache hit, appending suffix tokens\n\nThis achieves both correctness (100%) and performance (22.7x speedup on high prefix reuse workloads).<|im_end|>",
|
||||
];
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_cache_produces_identical_tokens() {
|
||||
// Get tokenizer path (download once, cached across tests)
|
||||
let tokenizer_path = match get_tokenizer_path().await {
|
||||
Some(path) => path,
|
||||
None => {
|
||||
println!("Skipping test - tokenizer not available");
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
// Create base tokenizer (no cache)
|
||||
let base_tokenizer = Arc::new(
|
||||
HuggingFaceTokenizer::from_file(tokenizer_path.to_str().unwrap())
|
||||
.expect("Failed to load base tokenizer"),
|
||||
);
|
||||
|
||||
// Create cached tokenizers with different configurations
|
||||
let l0_only_config = CacheConfig {
|
||||
enable_l0: true,
|
||||
l0_max_entries: 10_000,
|
||||
enable_l1: false,
|
||||
l1_max_memory: 0,
|
||||
};
|
||||
|
||||
let l1_only_config = CacheConfig {
|
||||
enable_l0: false,
|
||||
l0_max_entries: 0,
|
||||
enable_l1: true,
|
||||
l1_max_memory: 50 * 1024 * 1024,
|
||||
};
|
||||
|
||||
let l0_l1_config = CacheConfig {
|
||||
enable_l0: true,
|
||||
l0_max_entries: 10_000,
|
||||
enable_l1: true,
|
||||
l1_max_memory: 50 * 1024 * 1024,
|
||||
};
|
||||
|
||||
let l0_tokenizer = Arc::new(CachedTokenizer::new(base_tokenizer.clone(), l0_only_config));
|
||||
let l1_tokenizer = Arc::new(CachedTokenizer::new(base_tokenizer.clone(), l1_only_config));
|
||||
let l0_l1_tokenizer = Arc::new(CachedTokenizer::new(base_tokenizer.clone(), l0_l1_config));
|
||||
|
||||
println!(
|
||||
"\n=== Testing Cache Correctness Across {} Chat Turns ===\n",
|
||||
CHAT_TURNS.len()
|
||||
);
|
||||
|
||||
for (turn_idx, turn) in CHAT_TURNS.iter().enumerate() {
|
||||
println!("Turn {}: Testing {} chars", turn_idx + 1, turn.len());
|
||||
|
||||
// Tokenize with base (no cache)
|
||||
let base_encoding = base_tokenizer
|
||||
.encode(turn, false)
|
||||
.expect("Base tokenization failed");
|
||||
let base_tokens = base_encoding.token_ids();
|
||||
|
||||
// Tokenize with L0-only
|
||||
let l0_encoding = l0_tokenizer
|
||||
.encode(turn, false)
|
||||
.expect("L0 tokenization failed");
|
||||
let l0_tokens = l0_encoding.token_ids();
|
||||
|
||||
// Tokenize with L1-only
|
||||
let l1_encoding = l1_tokenizer
|
||||
.encode(turn, false)
|
||||
.expect("L1 tokenization failed");
|
||||
let l1_tokens = l1_encoding.token_ids();
|
||||
|
||||
// Tokenize with L0+L1
|
||||
let l0_l1_encoding = l0_l1_tokenizer
|
||||
.encode(turn, false)
|
||||
.expect("L0+L1 tokenization failed");
|
||||
let l0_l1_tokens = l0_l1_encoding.token_ids();
|
||||
|
||||
// Verify all configurations produce identical token IDs
|
||||
assert_eq!(
|
||||
base_tokens.len(),
|
||||
l0_tokens.len(),
|
||||
"Turn {}: L0 token count mismatch (base: {}, L0: {})",
|
||||
turn_idx + 1,
|
||||
base_tokens.len(),
|
||||
l0_tokens.len()
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
base_tokens.len(),
|
||||
l1_tokens.len(),
|
||||
"Turn {}: L1 token count mismatch (base: {}, L1: {})",
|
||||
turn_idx + 1,
|
||||
base_tokens.len(),
|
||||
l1_tokens.len()
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
base_tokens.len(),
|
||||
l0_l1_tokens.len(),
|
||||
"Turn {}: L0+L1 token count mismatch (base: {}, L0+L1: {})",
|
||||
turn_idx + 1,
|
||||
base_tokens.len(),
|
||||
l0_l1_tokens.len()
|
||||
);
|
||||
|
||||
// Compare token by token
|
||||
for (token_idx, (((base_token, l0_token), l1_token), l0_l1_token)) in base_tokens
|
||||
.iter()
|
||||
.zip(l0_tokens.iter())
|
||||
.zip(l1_tokens.iter())
|
||||
.zip(l0_l1_tokens.iter())
|
||||
.enumerate()
|
||||
{
|
||||
assert_eq!(
|
||||
base_token,
|
||||
l0_token,
|
||||
"Turn {}, token {}: L0 mismatch (base: {}, L0: {})",
|
||||
turn_idx + 1,
|
||||
token_idx,
|
||||
base_token,
|
||||
l0_token
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
base_token,
|
||||
l1_token,
|
||||
"Turn {}, token {}: L1 mismatch (base: {}, L1: {})",
|
||||
turn_idx + 1,
|
||||
token_idx,
|
||||
base_token,
|
||||
l1_token
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
base_token,
|
||||
l0_l1_token,
|
||||
"Turn {}, token {}: L0+L1 mismatch (base: {}, L0+L1: {})",
|
||||
turn_idx + 1,
|
||||
token_idx,
|
||||
base_token,
|
||||
l0_l1_token
|
||||
);
|
||||
}
|
||||
|
||||
println!(
|
||||
" ✓ All configurations produced identical {} tokens",
|
||||
base_tokens.len()
|
||||
);
|
||||
}
|
||||
|
||||
// Print cache statistics
|
||||
if let Some(l0_stats) = l0_tokenizer.cache_stats() {
|
||||
println!("\n=== L0 Cache Statistics ===");
|
||||
println!(" Hits: {}", l0_stats.hits);
|
||||
println!(" Misses: {}", l0_stats.misses);
|
||||
println!(
|
||||
" Hit rate: {:.2}%",
|
||||
if l0_stats.hits + l0_stats.misses > 0 {
|
||||
l0_stats.hits as f64 / (l0_stats.hits + l0_stats.misses) as f64 * 100.0
|
||||
} else {
|
||||
0.0
|
||||
}
|
||||
);
|
||||
println!(" Entries: {}", l0_stats.entries);
|
||||
}
|
||||
|
||||
if let Some(l1_stats) = l1_tokenizer.l1_cache_stats() {
|
||||
println!("\n=== L1 Cache Statistics ===");
|
||||
println!(" Hits: {}", l1_stats.hits);
|
||||
println!(" Misses: {}", l1_stats.misses);
|
||||
println!(
|
||||
" Hit rate: {:.2}%",
|
||||
if l1_stats.hits + l1_stats.misses > 0 {
|
||||
l1_stats.hits as f64 / (l1_stats.hits + l1_stats.misses) as f64 * 100.0
|
||||
} else {
|
||||
0.0
|
||||
}
|
||||
);
|
||||
println!(" Entries: {}", l1_stats.entries);
|
||||
println!(" Memory used: {} bytes", l1_stats.memory_bytes);
|
||||
}
|
||||
|
||||
if let Some(l0_stats) = l0_l1_tokenizer.cache_stats() {
|
||||
if let Some(l1_stats) = l0_l1_tokenizer.l1_cache_stats() {
|
||||
println!("\n=== L0+L1 Combined Cache Statistics ===");
|
||||
println!(" L0 Hits: {}", l0_stats.hits);
|
||||
println!(" L1 Hits: {}", l1_stats.hits);
|
||||
println!(
|
||||
" Total Hit rate: {:.2}%",
|
||||
if l0_stats.hits + l1_stats.hits + l0_stats.misses + l1_stats.misses > 0 {
|
||||
(l0_stats.hits + l1_stats.hits) as f64
|
||||
/ (l0_stats.hits + l1_stats.hits + l0_stats.misses + l1_stats.misses) as f64
|
||||
* 100.0
|
||||
} else {
|
||||
0.0
|
||||
}
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
println!("\n✓ All cache configurations produce identical tokenization results!");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_cache_correctness_with_edge_cases() {
|
||||
// Get tokenizer path (download once, cached across tests)
|
||||
let tokenizer_path = match get_tokenizer_path().await {
|
||||
Some(path) => path,
|
||||
None => {
|
||||
println!("Skipping test - tokenizer not available");
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
// Create base and cached tokenizers
|
||||
let base_tokenizer = Arc::new(
|
||||
HuggingFaceTokenizer::from_file(tokenizer_path.to_str().unwrap())
|
||||
.expect("Failed to load base tokenizer"),
|
||||
);
|
||||
|
||||
let cached_config = CacheConfig {
|
||||
enable_l0: true,
|
||||
l0_max_entries: 10_000,
|
||||
enable_l1: true,
|
||||
l1_max_memory: 50 * 1024 * 1024,
|
||||
};
|
||||
|
||||
let cached_tokenizer = Arc::new(CachedTokenizer::new(base_tokenizer.clone(), cached_config));
|
||||
|
||||
println!("\n=== Testing Edge Cases and Complex Patterns ===\n");
|
||||
|
||||
// Edge cases that stress-test the cache
|
||||
let edge_cases = [
|
||||
// Minimal messages
|
||||
("<|im_start|>system\n<|im_end|>", "Empty system message"),
|
||||
("<|im_start|>user\na<|im_end|>", "Single character"),
|
||||
|
||||
// Special token boundaries
|
||||
("<|im_start|>system\nA<|im_end|><|im_start|>user\nB<|im_end|><|im_start|>assistant\nC<|im_end|>", "Minimal multi-turn"),
|
||||
|
||||
// Repeated exact queries (L0 hit test)
|
||||
("<|im_start|>system\nYou are helpful.<|im_end|><|im_start|>user\nHello!<|im_end|>", "Repeated query 1"),
|
||||
("<|im_start|>system\nYou are helpful.<|im_end|><|im_start|>user\nHello!<|im_end|>", "Repeated query 2"),
|
||||
|
||||
// Same prefix, different suffix (L1 hit test)
|
||||
("<|im_start|>system\nYou are helpful.<|im_end|><|im_start|>user\nWhat is 1+1?<|im_end|>", "Same prefix, diff suffix 1"),
|
||||
("<|im_start|>system\nYou are helpful.<|im_end|><|im_start|>user\nWhat is 2+2?<|im_end|>", "Same prefix, diff suffix 2"),
|
||||
("<|im_start|>system\nYou are helpful.<|im_end|><|im_start|>user\nWhat is 3+3?<|im_end|>", "Same prefix, diff suffix 3"),
|
||||
|
||||
// Unicode stress tests
|
||||
("<|im_start|>system\n你好<|im_end|><|im_start|>user\n世界<|im_end|>", "Chinese characters"),
|
||||
("<|im_start|>system\nこんにちは<|im_end|><|im_start|>user\n世界<|im_end|>", "Japanese + Chinese"),
|
||||
("<|im_start|>system\n🚀💻🎉<|im_end|><|im_start|>user\n😀😃😄<|im_end|>", "Emoji only"),
|
||||
|
||||
// Whitespace edge cases
|
||||
("<|im_start|>system\n \n<|im_end|>", "Whitespace only"),
|
||||
("<|im_start|>system\n\n\n\n<|im_end|>", "Multiple newlines"),
|
||||
("<|im_start|>system\n\t\t\t<|im_end|>", "Tabs"),
|
||||
|
||||
// Long token sequences
|
||||
("<|im_start|>system\nThe quick brown fox jumps over the lazy dog. The quick brown fox jumps over the lazy dog. The quick brown fox jumps over the lazy dog.<|im_end|>", "Repeated phrase"),
|
||||
|
||||
// Special characters
|
||||
("<|im_start|>system\n!@#$%^&*()_+-=[]{}|;':\",./<>?<|im_end|>", "ASCII special chars"),
|
||||
("<|im_start|>system\n`~\\<|im_end|>", "Backtick and tilde"),
|
||||
|
||||
// Code with special formatting
|
||||
("<|im_start|>system\nCode: fn() -> Result<(), Box<dyn Error>><|im_end|>", "Rust generics"),
|
||||
("<|im_start|>system\nRegex: ^[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\\.[a-zA-Z]{2,}$<|im_end|>", "Email regex"),
|
||||
|
||||
// Very long single token sequences (testing buffer handling)
|
||||
("<|im_start|>system\naaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa<|im_end|>", "Repeated 'a'"),
|
||||
("<|im_start|>system\n0123456789012345678901234567890123456789012345678901234567890123456789012345678901234567890123456789<|im_end|>", "Repeated numbers"),
|
||||
];
|
||||
|
||||
let mut test_count = 0;
|
||||
let mut mismatch_count = 0;
|
||||
|
||||
for (query, description) in edge_cases.iter() {
|
||||
test_count += 1;
|
||||
|
||||
let base_tokens = base_tokenizer
|
||||
.encode(query, false)
|
||||
.expect("Base encoding failed")
|
||||
.token_ids()
|
||||
.to_vec();
|
||||
|
||||
let cached_tokens = cached_tokenizer
|
||||
.encode(query, false)
|
||||
.expect("Cached encoding failed")
|
||||
.token_ids()
|
||||
.to_vec();
|
||||
|
||||
if base_tokens != cached_tokens {
|
||||
mismatch_count += 1;
|
||||
println!(" ✗ {}: Token mismatch!", description);
|
||||
println!(
|
||||
" Base length: {}, Cached length: {}",
|
||||
base_tokens.len(),
|
||||
cached_tokens.len()
|
||||
);
|
||||
|
||||
// Show first few mismatching tokens for debugging
|
||||
for (i, (base, cached)) in base_tokens.iter().zip(cached_tokens.iter()).enumerate() {
|
||||
if base != cached {
|
||||
println!(" Token {}: base={}, cached={}", i, base, cached);
|
||||
if i >= 5 {
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
} else {
|
||||
println!(" ✓ {}: {} tokens", description, base_tokens.len());
|
||||
}
|
||||
}
|
||||
|
||||
assert_eq!(
|
||||
mismatch_count, 0,
|
||||
"{} out of {} edge cases failed!",
|
||||
mismatch_count, test_count
|
||||
);
|
||||
|
||||
// Print cache statistics
|
||||
if let Some(l0_stats) = cached_tokenizer.cache_stats() {
|
||||
println!("\n=== Cache Statistics ===");
|
||||
println!(
|
||||
" L0 Hits: {} ({:.1}% hit rate)",
|
||||
l0_stats.hits,
|
||||
if l0_stats.hits + l0_stats.misses > 0 {
|
||||
l0_stats.hits as f64 / (l0_stats.hits + l0_stats.misses) as f64 * 100.0
|
||||
} else {
|
||||
0.0
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
if let Some(l1_stats) = cached_tokenizer.l1_cache_stats() {
|
||||
println!(
|
||||
" L1 Hits: {} ({:.1}% hit rate)",
|
||||
l1_stats.hits,
|
||||
if l1_stats.hits + l1_stats.misses > 0 {
|
||||
l1_stats.hits as f64 / (l1_stats.hits + l1_stats.misses) as f64 * 100.0
|
||||
} else {
|
||||
0.0
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
println!("\n✓ All {} edge cases passed!", test_count);
|
||||
}
|
||||
@@ -1,570 +0,0 @@
|
||||
//! Integration tests for tokenizers using real tokenizer data
|
||||
//!
|
||||
//! These tests download the TinyLlama tokenizer from HuggingFace to verify our tokenizer
|
||||
//! implementation works correctly with real-world tokenizer files.
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use smg::tokenizer::{
|
||||
factory, huggingface::HuggingFaceTokenizer, sequence::Sequence, stop::*, stream::DecodeStream,
|
||||
traits::*,
|
||||
};
|
||||
|
||||
use crate::common::{ensure_tokenizer_cached, EXPECTED_HASHES, TEST_PROMPTS};
|
||||
|
||||
const LONG_TEST_PROMPTS: [(&str, &str); 6] = [
|
||||
("Tell me about the following text.", "Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat."),
|
||||
("Tell me about the following text.", "Duis aute irure dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat nulla pariatur. Excepteur sint occaecat cupidatat non proident, sunt in culpa qui officia deserunt mollit anim id est laborum."),
|
||||
("Tell me about the following text.", "Sed ut perspiciatis unde omnis iste natus error sit voluptatem accusantium doloremque laudantium, totam rem aperiam, eaque ipsa quae ab illo inventore veritatis et quasi architecto beatae vitae dicta sunt explicabo. Nemo enim ipsam voluptatem quia voluptas sit aspernatur aut odit aut fugit, sed quia consequuntur magni dolores eos qui ratione voluptatem sequi nesciunt."),
|
||||
("Tell me about the following text.", "Neque porro quisquam est, qui dolorem ipsum quia dolor sit amet, consectetur, adipisci velit, sed quia non numquam eius modi tempora incidunt ut labore et dolore magnam aliquam quaerat voluptatem."),
|
||||
// Tennis-themed prompt for variety
|
||||
("Tell me about the following text.", "In the ancient realm of Tennisia, the very magic of the land is drawn from the sport itself. Forehands light the skies, backhands carve the earth, and serves rumble like thunder across kingdoms. At the center of this balance lie four sacred Grand Slam relics: the Sapphire Trophy of Melbourne, the Emerald Chalice of Paris, the Ruby Crown of London, and the Diamond Orb of New York. Together, they keep the game's spirit alive.
|
||||
But the relics are scattered, guarded by champions of legendary skill. The first is the Fire King of Clay, ruler of the crimson courts, whose topspin arcs blaze high and heavy, scorching all who dare stand across from him. The second is the Tempest Trickster, master of the baseline fortress, whose footwork and precision can turn back any storm, and whose returns arrive as if pulled by invisible strings. The third is the Shadow-Dancer of the Highlands, a tactician who thrives in the long rallies of twilight, changing pace and spin until opponents lose their rhythm. The fourth and final guardian is a towering Diamond Titan, a net-charging colossus whose volleys shatter the air itself.
|
||||
Into this arena of gods steps the Silver-Wristed Knight — a player of impossible grace, whose game is an art form. His quest: to claim each relic not for glory, but to restore harmony to the rankings of the realm.
|
||||
He travels across the Kingdom of Clay, where the points stretch like marathons and the air tastes of iron; through the Grasslands of London, where the ball skids low and the margins are razor-thin; over the Hard Courts of the East, where rallies turn into duels of endurance; and finally to the Cathedral of Lights in New York, where night matches burn with fevered energy.
|
||||
Each battle is played under enchanted floodlights, the lines patrolled by spectral line judges whose calls are final. The crowd's roar swells with every break point, and the Silver-Wristed Knight's racket glows brightest when the match teeters at deuce. There are moments when doubt grips him — when his serve falters or his touch deserts him — but each challenge teaches a new stroke, culminating in the legendary Forehand of Dawn.
|
||||
When the last relic is claimed, he stands not as a conqueror but as a custodian of the game, knowing that rivalries forge the very magic he protects. The balance is restored — until the next season begins."),
|
||||
// Emoji stress test
|
||||
("Tell me about the following text.", "😀😃😄😁😆🥹😅😂🤣🥲☺️😊😇🙂🙃😉🤩😎 🤪🥳🤓🙄🤪😵👻")
|
||||
];
|
||||
|
||||
fn compute_hashes_for_tokenizer<E: Encoder>(tokenizer: &E, prompts: &[&str]) -> Vec<u64> {
|
||||
prompts
|
||||
.iter()
|
||||
.map(|&prompt| {
|
||||
tokenizer
|
||||
.encode(prompt, false)
|
||||
.expect("Failed to encode prompt")
|
||||
.get_hash()
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_huggingface_tokenizer_hashes() {
|
||||
let tokenizer_path = ensure_tokenizer_cached();
|
||||
let tokenizer = HuggingFaceTokenizer::from_file(tokenizer_path.to_str().unwrap())
|
||||
.expect("Failed to load HuggingFace tokenizer");
|
||||
|
||||
let prompt_hashes = compute_hashes_for_tokenizer(&tokenizer, &TEST_PROMPTS);
|
||||
|
||||
println!(
|
||||
"HF Tokenizer: {:?}\nComputed Hashes: {:?}\nExpected Hashes: {:?}",
|
||||
tokenizer_path, prompt_hashes, EXPECTED_HASHES
|
||||
);
|
||||
|
||||
assert_eq!(prompt_hashes, EXPECTED_HASHES);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_tokenizer_encode_decode_lifecycle() {
|
||||
let tokenizer_path = ensure_tokenizer_cached();
|
||||
let tokenizer = HuggingFaceTokenizer::from_file(tokenizer_path.to_str().unwrap())
|
||||
.expect("Failed to load HuggingFace tokenizer");
|
||||
|
||||
for prompt in TEST_PROMPTS.iter() {
|
||||
let encoding = tokenizer
|
||||
.encode(prompt, false)
|
||||
.expect("Failed to encode prompt");
|
||||
|
||||
let decoded = tokenizer
|
||||
.decode(encoding.token_ids(), false)
|
||||
.expect("Failed to decode token_ids");
|
||||
|
||||
assert_eq!(decoded, *prompt, "Encode-decode mismatch for: {}", prompt);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_sequence_operations() {
|
||||
let tokenizer_path = ensure_tokenizer_cached();
|
||||
let tokenizer = Arc::new(
|
||||
HuggingFaceTokenizer::from_file(tokenizer_path.to_str().unwrap())
|
||||
.expect("Failed to load tokenizer"),
|
||||
);
|
||||
|
||||
for prompt in TEST_PROMPTS.iter() {
|
||||
let encoding = tokenizer
|
||||
.encode(prompt, false)
|
||||
.expect("Failed to encode prompt");
|
||||
|
||||
let mut sequence = Sequence::new(tokenizer.clone());
|
||||
sequence
|
||||
.append_text(prompt, false)
|
||||
.expect("Failed to append text");
|
||||
|
||||
assert_eq!(
|
||||
sequence.len(),
|
||||
encoding.token_ids().len(),
|
||||
"Sequence length mismatch"
|
||||
);
|
||||
assert_eq!(sequence.text().unwrap(), *prompt, "Sequence text mismatch");
|
||||
|
||||
let mut decoder = Sequence::new(tokenizer.clone());
|
||||
let mut output = String::new();
|
||||
|
||||
for token_id in encoding.token_ids() {
|
||||
let text = decoder
|
||||
.append_token(*token_id)
|
||||
.expect("Failed to append token");
|
||||
output.push_str(&text);
|
||||
}
|
||||
|
||||
assert_eq!(decoder.len(), sequence.len(), "Decoder length mismatch");
|
||||
assert_eq!(
|
||||
decoder.token_ids(),
|
||||
sequence.token_ids(),
|
||||
"Token IDs mismatch"
|
||||
);
|
||||
assert_eq!(output, *prompt, "Incremental decode mismatch");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_decode_stream() {
|
||||
let tokenizer_path = ensure_tokenizer_cached();
|
||||
let tokenizer = Arc::new(
|
||||
HuggingFaceTokenizer::from_file(tokenizer_path.to_str().unwrap())
|
||||
.expect("Failed to load tokenizer"),
|
||||
);
|
||||
|
||||
for prompt in TEST_PROMPTS.iter() {
|
||||
let encoding = tokenizer
|
||||
.encode(prompt, false)
|
||||
.expect("Failed to encode prompt");
|
||||
|
||||
let mut decoder = DecodeStream::new(tokenizer.clone(), &[], false);
|
||||
let mut output = String::new();
|
||||
|
||||
for token_id in encoding.token_ids() {
|
||||
if let Some(text) = decoder.step(*token_id).expect("Failed to decode token") {
|
||||
output.push_str(&text);
|
||||
}
|
||||
}
|
||||
|
||||
assert_eq!(output, *prompt, "DecodeStream output mismatch");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_long_sequence_incremental_decode_with_prefill() {
|
||||
let tokenizer_path = ensure_tokenizer_cached();
|
||||
let tokenizer = Arc::new(
|
||||
HuggingFaceTokenizer::from_file(tokenizer_path.to_str().unwrap())
|
||||
.expect("Failed to load tokenizer"),
|
||||
);
|
||||
|
||||
for (input_text, output_text) in LONG_TEST_PROMPTS.iter() {
|
||||
let input_encoding = tokenizer
|
||||
.encode(input_text, false)
|
||||
.expect("Failed to encode input");
|
||||
|
||||
let output_encoding = tokenizer
|
||||
.encode(output_text, false)
|
||||
.expect("Failed to encode output");
|
||||
|
||||
let mut decoder = DecodeStream::new(tokenizer.clone(), input_encoding.token_ids(), false);
|
||||
|
||||
let mut output = String::new();
|
||||
for token_id in output_encoding.token_ids() {
|
||||
if let Some(text) = decoder.step(*token_id).expect("Failed to decode token") {
|
||||
output.push_str(&text);
|
||||
}
|
||||
}
|
||||
|
||||
assert_eq!(output.trim(), *output_text, "Long sequence decode mismatch");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_stop_sequence_decoder() {
|
||||
let tokenizer_path = ensure_tokenizer_cached();
|
||||
let tokenizer = Arc::new(
|
||||
HuggingFaceTokenizer::from_file(tokenizer_path.to_str().unwrap())
|
||||
.expect("Failed to load tokenizer"),
|
||||
);
|
||||
|
||||
let test_cases = vec![
|
||||
(
|
||||
"Hello world! Stop here. Continue after.",
|
||||
"Stop",
|
||||
"Hello world! ",
|
||||
),
|
||||
("Testing stop sequences.", ".", "Testing stop sequences"),
|
||||
("No stop sequence here", "xyz", "No stop sequence here"),
|
||||
];
|
||||
|
||||
for (input, stop_seq, expected) in test_cases {
|
||||
let config = StopSequenceConfig::default().with_stop_sequence(stop_seq);
|
||||
|
||||
let mut decoder = StopSequenceDecoder::new(tokenizer.clone(), config, false);
|
||||
|
||||
let encoding = tokenizer.encode(input, false).expect("Failed to encode");
|
||||
let mut output = String::new();
|
||||
let mut stopped = false;
|
||||
|
||||
for token_id in encoding.token_ids() {
|
||||
match decoder.process_token(*token_id).unwrap() {
|
||||
SequenceDecoderOutput::Text(text) => output.push_str(&text),
|
||||
SequenceDecoderOutput::StoppedWithText(text) => {
|
||||
output.push_str(&text);
|
||||
stopped = true;
|
||||
break;
|
||||
}
|
||||
SequenceDecoderOutput::Stopped => {
|
||||
stopped = true;
|
||||
break;
|
||||
}
|
||||
SequenceDecoderOutput::Held => {}
|
||||
}
|
||||
}
|
||||
|
||||
if !stopped {
|
||||
// Flush any remaining text
|
||||
if let SequenceDecoderOutput::Text(text) = decoder.flush() {
|
||||
output.push_str(&text);
|
||||
}
|
||||
}
|
||||
|
||||
println!(
|
||||
"Input: '{}', Stop: '{}', Output: '{}', Expected: '{}'",
|
||||
input, stop_seq, output, expected
|
||||
);
|
||||
|
||||
// The test should check if output starts with expected
|
||||
// since stop sequences might not be perfectly aligned with token boundaries
|
||||
assert!(
|
||||
output.starts_with(expected) || output == input,
|
||||
"Stop sequence test failed"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_factory_creation() {
|
||||
let tokenizer_path = ensure_tokenizer_cached();
|
||||
let tokenizer = factory::create_tokenizer(tokenizer_path.to_str().unwrap())
|
||||
.expect("Failed to create tokenizer via factory");
|
||||
|
||||
let encoding = tokenizer
|
||||
.encode(TEST_PROMPTS[0], false)
|
||||
.expect("Failed to encode");
|
||||
|
||||
let decoded = tokenizer
|
||||
.decode(encoding.token_ids(), false)
|
||||
.expect("Failed to decode");
|
||||
|
||||
assert_eq!(decoded, TEST_PROMPTS[0]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_batch_encoding() {
|
||||
let tokenizer_path = ensure_tokenizer_cached();
|
||||
let tokenizer = HuggingFaceTokenizer::from_file(tokenizer_path.to_str().unwrap())
|
||||
.expect("Failed to load tokenizer");
|
||||
|
||||
let encodings = tokenizer
|
||||
.encode_batch(&TEST_PROMPTS, false)
|
||||
.expect("Failed to batch encode");
|
||||
|
||||
assert_eq!(encodings.len(), TEST_PROMPTS.len());
|
||||
|
||||
for (i, encoding) in encodings.iter().enumerate() {
|
||||
let decoded = tokenizer
|
||||
.decode(encoding.token_ids(), false)
|
||||
.expect("Failed to decode");
|
||||
assert_eq!(decoded, TEST_PROMPTS[i]);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_special_tokens() {
|
||||
use smg::tokenizer::traits::Tokenizer as TokenizerTrait;
|
||||
|
||||
let tokenizer_path = ensure_tokenizer_cached();
|
||||
let tokenizer = HuggingFaceTokenizer::from_file(tokenizer_path.to_str().unwrap())
|
||||
.expect("Failed to load tokenizer");
|
||||
|
||||
let special_tokens = tokenizer.get_special_tokens();
|
||||
|
||||
// TinyLlama should have at least BOS and EOS tokens
|
||||
assert!(special_tokens.bos_token.is_some());
|
||||
assert!(special_tokens.eos_token.is_some());
|
||||
|
||||
println!("Special tokens: {:?}", special_tokens);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_thread_safety() {
|
||||
use std::thread;
|
||||
|
||||
let tokenizer_path = ensure_tokenizer_cached();
|
||||
let tokenizer = Arc::new(
|
||||
HuggingFaceTokenizer::from_file(tokenizer_path.to_str().unwrap())
|
||||
.expect("Failed to load tokenizer"),
|
||||
);
|
||||
|
||||
let handles: Vec<_> = TEST_PROMPTS
|
||||
.iter()
|
||||
.map(|&prompt| {
|
||||
let tokenizer_clone = tokenizer.clone();
|
||||
thread::spawn(move || {
|
||||
let encoding = tokenizer_clone
|
||||
.encode(prompt, false)
|
||||
.expect("Failed to encode in thread");
|
||||
let decoded = tokenizer_clone
|
||||
.decode(encoding.token_ids(), false)
|
||||
.expect("Failed to decode in thread");
|
||||
assert_eq!(decoded, prompt);
|
||||
})
|
||||
})
|
||||
.collect();
|
||||
|
||||
for handle in handles {
|
||||
handle.join().expect("Thread panicked");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_chat_template_discovery() {
|
||||
use std::fs;
|
||||
|
||||
use tempfile::TempDir;
|
||||
|
||||
// Create a temporary directory with test files
|
||||
let temp_dir = TempDir::new().expect("Failed to create temp dir");
|
||||
let dir_path = temp_dir.path();
|
||||
|
||||
// Copy a real tokenizer.json file for testing
|
||||
// We'll use the TinyLlama tokenizer that's already cached
|
||||
let cached_tokenizer = ensure_tokenizer_cached();
|
||||
let tokenizer_path = dir_path.join("tokenizer.json");
|
||||
fs::copy(&cached_tokenizer, &tokenizer_path).expect("Failed to copy tokenizer file");
|
||||
|
||||
// Test 1: With chat_template.jinja file
|
||||
let jinja_path = dir_path.join("chat_template.jinja");
|
||||
fs::write(&jinja_path, "{{ messages }}").expect("Failed to write chat template");
|
||||
|
||||
let tokenizer = HuggingFaceTokenizer::from_file(tokenizer_path.to_str().unwrap());
|
||||
assert!(
|
||||
tokenizer.is_ok(),
|
||||
"Should load tokenizer with chat template"
|
||||
);
|
||||
|
||||
// Clean up for next test
|
||||
fs::remove_file(&jinja_path).ok();
|
||||
|
||||
// Test 2: With tokenizer_config.json containing chat_template
|
||||
let config_path = dir_path.join("tokenizer_config.json");
|
||||
fs::write(&config_path, r#"{"chat_template": "{{ messages }}"}"#)
|
||||
.expect("Failed to write config");
|
||||
|
||||
let tokenizer = HuggingFaceTokenizer::from_file(tokenizer_path.to_str().unwrap());
|
||||
assert!(
|
||||
tokenizer.is_ok(),
|
||||
"Should load tokenizer with embedded template"
|
||||
);
|
||||
|
||||
// Test 3: No chat template
|
||||
fs::remove_file(&config_path).ok();
|
||||
let tokenizer = HuggingFaceTokenizer::from_file(tokenizer_path.to_str().unwrap());
|
||||
assert!(
|
||||
tokenizer.is_ok(),
|
||||
"Should load tokenizer without chat template"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_load_chat_template_from_local_file() {
|
||||
use std::fs;
|
||||
|
||||
use tempfile::TempDir;
|
||||
|
||||
// Test 1: Load tokenizer with explicit chat template path
|
||||
let temp_dir = TempDir::new().expect("Failed to create temp dir");
|
||||
let dir_path = temp_dir.path();
|
||||
|
||||
// Copy a real tokenizer for testing
|
||||
let cached_tokenizer = ensure_tokenizer_cached();
|
||||
let tokenizer_path = dir_path.join("tokenizer.json");
|
||||
fs::copy(&cached_tokenizer, &tokenizer_path).expect("Failed to copy tokenizer");
|
||||
|
||||
// Create a chat template file
|
||||
let template_path = dir_path.join("my_template.jinja");
|
||||
let template_content = r#"{% for message in messages %}{{ message.role }}: {{ message.content }}
|
||||
{% endfor %}"#;
|
||||
fs::write(&template_path, template_content).expect("Failed to write template");
|
||||
|
||||
// Load tokenizer with explicit template path
|
||||
let tokenizer = HuggingFaceTokenizer::from_file_with_chat_template(
|
||||
tokenizer_path.to_str().unwrap(),
|
||||
Some(template_path.to_str().unwrap()),
|
||||
);
|
||||
assert!(
|
||||
tokenizer.is_ok(),
|
||||
"Should load tokenizer with explicit template path"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_tinyllama_embedded_template() {
|
||||
use smg::tokenizer::hub::download_tokenizer_from_hf;
|
||||
|
||||
// Skip in CI without HF_TOKEN
|
||||
|
||||
// Test 2: TinyLlama has chat template embedded in tokenizer_config.json
|
||||
match download_tokenizer_from_hf("TinyLlama/TinyLlama-1.1B-Chat-v1.0").await {
|
||||
Ok(cache_dir) => {
|
||||
// Verify tokenizer_config.json exists
|
||||
let config_path = cache_dir.join("tokenizer_config.json");
|
||||
assert!(config_path.exists(), "tokenizer_config.json should exist");
|
||||
|
||||
// Load the config and check for chat_template
|
||||
let config_content =
|
||||
std::fs::read_to_string(&config_path).expect("Failed to read config");
|
||||
assert!(
|
||||
config_content.contains("\"chat_template\""),
|
||||
"TinyLlama should have embedded chat_template in config"
|
||||
);
|
||||
|
||||
// Load tokenizer and verify it has chat template
|
||||
let tokenizer_path = cache_dir.join("tokenizer.json");
|
||||
let _tokenizer = HuggingFaceTokenizer::from_file(tokenizer_path.to_str().unwrap())
|
||||
.expect("Failed to load tokenizer");
|
||||
|
||||
println!(
|
||||
"✓ TinyLlama: Loaded tokenizer with embedded template from tokenizer_config.json"
|
||||
);
|
||||
}
|
||||
Err(e) => {
|
||||
println!("Download test skipped due to error: {}", e);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_qwen3_next_embedded_template() {
|
||||
use smg::tokenizer::hub::download_tokenizer_from_hf;
|
||||
|
||||
// Test 3: Qwen3-Next has chat template in tokenizer_config.json
|
||||
match download_tokenizer_from_hf("Qwen/Qwen3-Next-80B-A3B-Instruct").await {
|
||||
Ok(cache_dir) => {
|
||||
let config_path = cache_dir.join("tokenizer_config.json");
|
||||
assert!(config_path.exists(), "tokenizer_config.json should exist");
|
||||
|
||||
// Verify chat_template in config
|
||||
let config_content =
|
||||
std::fs::read_to_string(&config_path).expect("Failed to read config");
|
||||
assert!(
|
||||
config_content.contains("\"chat_template\""),
|
||||
"Qwen3-Next should have chat_template in tokenizer_config.json"
|
||||
);
|
||||
|
||||
// Load tokenizer
|
||||
let tokenizer_path = cache_dir.join("tokenizer.json");
|
||||
if tokenizer_path.exists() {
|
||||
let _tokenizer = HuggingFaceTokenizer::from_file(tokenizer_path.to_str().unwrap())
|
||||
.expect("Failed to load tokenizer");
|
||||
println!("✓ Qwen3-Next: Loaded tokenizer with embedded template");
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
println!("Download test skipped due to error: {}", e);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_qwen3_vl_json_template_priority() {
|
||||
use smg::tokenizer::hub::download_tokenizer_from_hf;
|
||||
|
||||
// Test 4: Qwen3-VL has both tokenizer_config.json template and chat_template.json
|
||||
// Should prioritize chat_template.json
|
||||
match download_tokenizer_from_hf("Qwen/Qwen3-VL-235B-A22B-Instruct").await {
|
||||
Ok(cache_dir) => {
|
||||
// Check for chat_template.json
|
||||
let json_template_path = cache_dir.join("chat_template.json");
|
||||
let has_json_template = json_template_path.exists();
|
||||
|
||||
// Also check tokenizer_config.json
|
||||
let config_path = cache_dir.join("tokenizer_config.json");
|
||||
assert!(config_path.exists(), "tokenizer_config.json should exist");
|
||||
|
||||
if has_json_template {
|
||||
let json_content = std::fs::read_to_string(&json_template_path)
|
||||
.expect("Failed to read chat_template.json");
|
||||
println!("✓ Qwen3-VL: Found chat_template.json (should be prioritized)");
|
||||
|
||||
// Verify it contains jinja template
|
||||
assert!(
|
||||
!json_content.is_empty(),
|
||||
"chat_template.json should contain template"
|
||||
);
|
||||
}
|
||||
|
||||
// Load tokenizer - it should use the appropriate template
|
||||
let tokenizer_path = cache_dir.join("tokenizer.json");
|
||||
if tokenizer_path.exists() {
|
||||
let _tokenizer = HuggingFaceTokenizer::from_file(tokenizer_path.to_str().unwrap())
|
||||
.expect("Failed to load tokenizer");
|
||||
println!("✓ Qwen3-VL: Loaded tokenizer with template priority handling");
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
println!("Download test skipped due to error: {}", e);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_llava_separate_jinja_template() {
|
||||
use smg::tokenizer::hub::download_tokenizer_from_hf;
|
||||
|
||||
// Test 5: llava has chat_template.jinja as a separate file, not in tokenizer_config.json
|
||||
match download_tokenizer_from_hf("llava-hf/llava-1.5-7b-hf").await {
|
||||
Ok(cache_dir) => {
|
||||
// Check for .jinja file
|
||||
let jinja_path = cache_dir.join("chat_template.jinja");
|
||||
let has_jinja = jinja_path.exists()
|
||||
|| std::fs::read_dir(&cache_dir)
|
||||
.map(|entries| {
|
||||
entries.filter_map(|e| e.ok()).any(|e| {
|
||||
e.file_name()
|
||||
.to_str()
|
||||
.is_some_and(|name| name.ends_with(".jinja"))
|
||||
})
|
||||
})
|
||||
.unwrap_or(false);
|
||||
|
||||
if has_jinja {
|
||||
println!("✓ llava: Found separate .jinja chat template file");
|
||||
}
|
||||
|
||||
// Check tokenizer_config.json - should NOT have embedded template
|
||||
let config_path = cache_dir.join("tokenizer_config.json");
|
||||
if config_path.exists() {
|
||||
let config_content =
|
||||
std::fs::read_to_string(&config_path).expect("Failed to read config");
|
||||
|
||||
// llava might not have chat_template in config
|
||||
if !config_content.contains("\"chat_template\"") {
|
||||
println!("✓ llava: No embedded template in config (as expected)");
|
||||
}
|
||||
}
|
||||
|
||||
// Load tokenizer - should auto-discover the .jinja file
|
||||
let tokenizer_path = cache_dir.join("tokenizer.json");
|
||||
if tokenizer_path.exists() {
|
||||
let tokenizer = HuggingFaceTokenizer::from_file(tokenizer_path.to_str().unwrap());
|
||||
if tokenizer.is_ok() {
|
||||
println!("✓ llava: Loaded tokenizer with auto-discovered .jinja template");
|
||||
} else {
|
||||
println!("Note: llava tokenizer loading failed - might need specific handling");
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
println!("Download test skipped due to error: {}", e);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,6 +0,0 @@
|
||||
//! Tokenizer integration tests
|
||||
|
||||
#[path = "common/mod.rs"]
|
||||
pub mod common;
|
||||
|
||||
mod tokenizer;
|
||||
Reference in New Issue
Block a user