149 lines
4.1 KiB
Rust
149 lines
4.1 KiB
Rust
//! Advanced tokenization with BPE/SentencePiece trainers and multimodal support
|
|
//!
|
|
//! This crate provides comprehensive tokenization capabilities including:
|
|
//! - Advanced Byte-Pair Encoding (BPE) with trainers
|
|
//! - `SentencePiece` tokenization
|
|
//! - Multimodal tokenization for images and audio
|
|
//! - CJK/RTL language support
|
|
//! - Dynamic vocabulary management
|
|
//! - Complete encoding/decoding pipelines
|
|
|
|
#![deny(missing_docs)]
|
|
|
|
pub mod bpe;
|
|
pub mod multimodal;
|
|
pub mod pipeline;
|
|
pub mod sentencepiece;
|
|
pub mod trainer;
|
|
pub mod vocabulary;
|
|
|
|
// Re-exports for convenience
|
|
pub use bpe::{BpeConfig, BpeTokenizer, BpeTrainer};
|
|
pub use multimodal::{
|
|
AudioFrame, AudioTokenizationConfig, AudioTokenizationMethod, ImagePatch,
|
|
ImageTokenizationConfig, ImageTokenizationMethod, ModalityType, MultimodalConfig,
|
|
MultimodalTokenizer,
|
|
};
|
|
pub use pipeline::{PipelineConfig, TokenizationPipeline};
|
|
pub use sentencepiece::{
|
|
SentencePieceConfig, SentencePieceModelType, SentencePieceTokenizer, SentencePieceTrainer,
|
|
SubwordPiece,
|
|
};
|
|
pub use trainer::{TrainedTokenizer, Trainer, TrainingAlgorithm, TrainingConfig, TrainingData};
|
|
pub use vocabulary::{Vocabulary, VocabularyBuilder};
|
|
|
|
/// Common error types for tokenization operations
|
|
#[derive(Debug, thiserror::Error)]
|
|
pub enum TokenizationError {
|
|
/// Invalid input data
|
|
#[error("Invalid input: {0}")]
|
|
InvalidInput(String),
|
|
|
|
/// Vocabulary not found
|
|
#[error("Vocabulary error: {0}")]
|
|
VocabularyError(String),
|
|
|
|
/// Training error
|
|
#[error("Training error: {0}")]
|
|
TrainingError(String),
|
|
|
|
/// IO error
|
|
#[error("IO error: {0}")]
|
|
IoError(#[from] std::io::Error),
|
|
|
|
/// Serialization error
|
|
#[error("Serialization error: {0}")]
|
|
SerializationError(#[from] serde_json::Error),
|
|
|
|
/// Unicode handling error
|
|
#[error("Unicode error: {0}")]
|
|
UnicodeError(String),
|
|
|
|
/// Multimodal processing error
|
|
#[error("Multimodal error: {0}")]
|
|
MultimodalError(String),
|
|
}
|
|
|
|
/// Result type for tokenization operations
|
|
pub type Result<T> = std::result::Result<T, TokenizationError>;
|
|
|
|
/// Token ID type
|
|
pub type TokenId = u32;
|
|
|
|
/// A token with its ID and text representation
|
|
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
|
|
pub struct Token {
|
|
/// Unique identifier for the token
|
|
pub id: TokenId,
|
|
/// Text representation of the token
|
|
pub text: String,
|
|
/// Optional frequency information
|
|
pub frequency: Option<u64>,
|
|
}
|
|
|
|
impl Token {
|
|
/// Create a new token
|
|
#[must_use]
|
|
pub fn new(id: TokenId, text: String, frequency: Option<u64>) -> Self {
|
|
Self {
|
|
id,
|
|
text,
|
|
frequency,
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Tokenization statistics
|
|
#[derive(Debug, Clone, Default, serde::Serialize, serde::Deserialize)]
|
|
pub struct TokenizationStats {
|
|
/// Number of tokens processed
|
|
pub token_count: usize,
|
|
/// Number of unique tokens
|
|
pub unique_tokens: usize,
|
|
/// Average token length
|
|
pub avg_token_length: f64,
|
|
/// Processing time in milliseconds
|
|
pub processing_time_ms: u64,
|
|
}
|
|
|
|
/// Common trait for all tokenizers
|
|
#[async_trait::async_trait]
|
|
pub trait Tokenizer: Send + Sync + std::fmt::Debug {
|
|
/// Encode text into token IDs
|
|
async fn encode(&self, text: &str) -> Result<Vec<TokenId>>;
|
|
|
|
/// Decode token IDs back to text
|
|
async fn decode(&self, token_ids: &[TokenId]) -> Result<String>;
|
|
|
|
/// Get vocabulary size
|
|
fn vocab_size(&self) -> usize;
|
|
|
|
/// Get tokenization statistics
|
|
fn get_stats(&self) -> TokenizationStats;
|
|
|
|
/// Check if tokenizer supports the given text
|
|
fn supports(&self, text: &str) -> bool;
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn token_creation() {
|
|
let token = Token::new(42, "hello".to_string(), Some(100));
|
|
assert_eq!(token.id, 42);
|
|
assert_eq!(token.text, "hello");
|
|
assert_eq!(token.frequency, Some(100));
|
|
}
|
|
|
|
#[test]
|
|
fn tokenization_stats_default() {
|
|
let stats = TokenizationStats::default();
|
|
assert_eq!(stats.token_count, 0);
|
|
assert_eq!(stats.unique_tokens, 0);
|
|
assert_eq!(stats.avg_token_length, 0.0);
|
|
assert_eq!(stats.processing_time_ms, 0);
|
|
}
|
|
}
|