395 lines
12 KiB
Rust
395 lines
12 KiB
Rust
//! Universal tokenization trainer supporting multiple algorithms
|
|
//!
|
|
//! This module provides a unified training interface for different tokenization
|
|
//! algorithms including BPE, `SentencePiece` unigram, and vocabulary-based approaches.
|
|
|
|
use crate::{
|
|
Result, TokenizationError, Tokenizer,
|
|
bpe::{BpeConfig, BpeTokenizer, BpeTrainer},
|
|
multimodal::{MultimodalConfig, MultimodalTokenizer},
|
|
sentencepiece::{SentencePieceConfig, SentencePieceTokenizer},
|
|
};
|
|
use serde::{Deserialize, Serialize};
|
|
use std::path::Path;
|
|
|
|
/// Training algorithm types
|
|
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
|
pub enum TrainingAlgorithm {
|
|
/// Byte-Pair Encoding
|
|
Bpe(BpeConfig),
|
|
/// `SentencePiece` Unigram
|
|
SentencePieceUnigram(SentencePieceConfig),
|
|
/// `SentencePiece` BPE
|
|
SentencePieceBpe(SentencePieceConfig),
|
|
/// Multimodal tokenization
|
|
Multimodal(MultimodalConfig),
|
|
}
|
|
|
|
/// Training data source
|
|
#[derive(Debug, Clone)]
|
|
pub enum TrainingData {
|
|
/// Text data from string
|
|
Text(String),
|
|
/// Text data from file path
|
|
TextFile(String),
|
|
/// Image data for multimodal training
|
|
Images(Vec<Vec<u8>>),
|
|
/// Audio data for multimodal training
|
|
Audio(Vec<Vec<f32>>),
|
|
/// Mixed multimodal data
|
|
Multimodal {
|
|
/// Text training data
|
|
text: Vec<String>,
|
|
/// Image training data (raw bytes)
|
|
images: Vec<Vec<u8>>,
|
|
/// Audio training data (waveform samples)
|
|
audio: Vec<Vec<f32>>,
|
|
},
|
|
}
|
|
|
|
/// Training configuration
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct TrainingConfig {
|
|
/// Training algorithm and its configuration
|
|
pub algorithm: TrainingAlgorithm,
|
|
/// Number of training iterations
|
|
pub max_iterations: usize,
|
|
/// Early stopping patience
|
|
pub patience: usize,
|
|
/// Validation split ratio
|
|
pub validation_split: f32,
|
|
/// Random seed for reproducibility
|
|
pub seed: Option<u64>,
|
|
/// Enable progress logging
|
|
pub verbose: bool,
|
|
/// Save intermediate models
|
|
pub save_checkpoints: bool,
|
|
}
|
|
|
|
impl Default for TrainingConfig {
|
|
fn default() -> Self {
|
|
Self {
|
|
algorithm: TrainingAlgorithm::Bpe(BpeConfig::default()),
|
|
max_iterations: 1000,
|
|
patience: 10,
|
|
validation_split: 0.1,
|
|
seed: None,
|
|
verbose: true,
|
|
save_checkpoints: false,
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Training metrics and statistics
|
|
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
|
|
pub struct TrainingMetrics {
|
|
/// Training loss history
|
|
pub loss_history: Vec<f64>,
|
|
/// Validation loss history
|
|
pub validation_loss: Vec<f64>,
|
|
/// Vocabulary growth over iterations
|
|
pub vocab_growth: Vec<usize>,
|
|
/// Training time in seconds
|
|
pub training_time_seconds: f64,
|
|
/// Final vocabulary size
|
|
pub final_vocab_size: usize,
|
|
/// Number of iterations completed
|
|
pub iterations_completed: usize,
|
|
/// Whether training converged
|
|
pub converged: bool,
|
|
}
|
|
|
|
/// Universal tokenization trainer
|
|
pub struct Trainer {
|
|
config: TrainingConfig,
|
|
metrics: TrainingMetrics,
|
|
}
|
|
|
|
impl Trainer {
|
|
/// Create a new trainer with configuration
|
|
#[must_use]
|
|
pub fn new(config: TrainingConfig) -> Self {
|
|
// Set random seed if provided
|
|
if let Some(seed) = config.seed {
|
|
use rand::{SeedableRng, rngs::StdRng};
|
|
let _rng = StdRng::seed_from_u64(seed);
|
|
}
|
|
|
|
Self {
|
|
config,
|
|
metrics: TrainingMetrics::default(),
|
|
}
|
|
}
|
|
|
|
/// Train a tokenizer with the provided data
|
|
pub async fn train(&mut self, data: TrainingData) -> Result<TrainedTokenizer> {
|
|
let start_time = std::time::Instant::now();
|
|
|
|
if self.config.verbose {
|
|
tracing::info!(
|
|
"Starting tokenization training with algorithm: {:?}",
|
|
std::mem::discriminant(&self.config.algorithm)
|
|
);
|
|
}
|
|
|
|
let tokenizer = match &self.config.algorithm {
|
|
TrainingAlgorithm::Bpe(bpe_config) => self.train_bpe(data, bpe_config.clone()).await?,
|
|
TrainingAlgorithm::SentencePieceUnigram(sp_config) => {
|
|
self.train_sentencepiece_unigram(data, sp_config.clone())
|
|
.await?
|
|
}
|
|
TrainingAlgorithm::SentencePieceBpe(sp_config) => {
|
|
self.train_sentencepiece_bpe(data, sp_config.clone())
|
|
.await?
|
|
}
|
|
TrainingAlgorithm::Multimodal(mm_config) => {
|
|
self.train_multimodal(data, mm_config.clone()).await?
|
|
}
|
|
};
|
|
|
|
self.metrics.training_time_seconds = start_time.elapsed().as_secs_f64();
|
|
self.metrics.converged = true;
|
|
|
|
if self.config.verbose {
|
|
tracing::info!(
|
|
"Training completed in {:.2}s with final vocab size: {}",
|
|
self.metrics.training_time_seconds,
|
|
self.metrics.final_vocab_size
|
|
);
|
|
}
|
|
|
|
Ok(tokenizer)
|
|
}
|
|
|
|
/// Train BPE tokenizer
|
|
async fn train_bpe(
|
|
&mut self,
|
|
data: TrainingData,
|
|
config: BpeConfig,
|
|
) -> Result<TrainedTokenizer> {
|
|
let mut trainer = BpeTrainer::new(config);
|
|
|
|
let texts = self.extract_text_data(data)?;
|
|
|
|
// Add training data with progress tracking
|
|
for (i, text) in texts.iter().enumerate() {
|
|
trainer.add_text(text)?;
|
|
|
|
if self.config.verbose && i % 1000 == 0 {
|
|
tracing::debug!("Processed {} texts", i);
|
|
}
|
|
}
|
|
|
|
let tokenizer = trainer.train()?;
|
|
self.metrics.final_vocab_size = tokenizer.vocab_size();
|
|
|
|
Ok(TrainedTokenizer::Bpe(tokenizer))
|
|
}
|
|
|
|
/// Train `SentencePiece` Unigram tokenizer
|
|
async fn train_sentencepiece_unigram(
|
|
&mut self,
|
|
_data: TrainingData,
|
|
config: SentencePieceConfig,
|
|
) -> Result<TrainedTokenizer> {
|
|
// This would require a full SentencePiece implementation
|
|
// For now, return a basic SentencePiece tokenizer
|
|
let tokenizer = SentencePieceTokenizer::new(config);
|
|
self.metrics.final_vocab_size = tokenizer.vocab_size();
|
|
|
|
Ok(TrainedTokenizer::SentencePiece(tokenizer))
|
|
}
|
|
|
|
/// Train `SentencePiece` BPE tokenizer
|
|
async fn train_sentencepiece_bpe(
|
|
&mut self,
|
|
_data: TrainingData,
|
|
config: SentencePieceConfig,
|
|
) -> Result<TrainedTokenizer> {
|
|
let tokenizer = SentencePieceTokenizer::new(config);
|
|
self.metrics.final_vocab_size = tokenizer.vocab_size();
|
|
|
|
Ok(TrainedTokenizer::SentencePiece(tokenizer))
|
|
}
|
|
|
|
/// Train multimodal tokenizer
|
|
async fn train_multimodal(
|
|
&mut self,
|
|
_data: TrainingData,
|
|
config: MultimodalConfig,
|
|
) -> Result<TrainedTokenizer> {
|
|
let tokenizer = MultimodalTokenizer::new(config);
|
|
self.metrics.final_vocab_size = tokenizer.vocab_size();
|
|
|
|
Ok(TrainedTokenizer::Multimodal(tokenizer))
|
|
}
|
|
|
|
/// Extract text data from various data sources
|
|
fn extract_text_data(&self, data: TrainingData) -> Result<Vec<String>> {
|
|
match data {
|
|
TrainingData::Text(text) => Ok(vec![text]),
|
|
TrainingData::TextFile(path) => {
|
|
let content = std::fs::read_to_string(&path).map_err(TokenizationError::IoError)?;
|
|
Ok(vec![content])
|
|
}
|
|
TrainingData::Images(_) => Err(TokenizationError::InvalidInput(
|
|
"Cannot extract text from image data".to_string(),
|
|
)),
|
|
TrainingData::Audio(_) => Err(TokenizationError::InvalidInput(
|
|
"Cannot extract text from audio data".to_string(),
|
|
)),
|
|
TrainingData::Multimodal { text, .. } => Ok(text),
|
|
}
|
|
}
|
|
|
|
/// Get current training metrics
|
|
#[must_use]
|
|
pub fn metrics(&self) -> &TrainingMetrics {
|
|
&self.metrics
|
|
}
|
|
|
|
/// Save trained model to file
|
|
pub async fn save_model<P: AsRef<Path>>(
|
|
&self,
|
|
tokenizer: &TrainedTokenizer,
|
|
path: P,
|
|
) -> Result<()> {
|
|
let serialized = bincode::serialize(tokenizer)
|
|
.map_err(|e| TokenizationError::InvalidInput(format!("Serialization error: {e}")))?;
|
|
|
|
tokio::fs::write(path, serialized)
|
|
.await
|
|
.map_err(TokenizationError::IoError)?;
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Load trained model from file
|
|
pub async fn load_model<P: AsRef<Path>>(path: P) -> Result<TrainedTokenizer> {
|
|
let data = tokio::fs::read(path)
|
|
.await
|
|
.map_err(TokenizationError::IoError)?;
|
|
|
|
let tokenizer = bincode::deserialize(&data)
|
|
.map_err(|e| TokenizationError::InvalidInput(format!("Deserialization error: {e}")))?;
|
|
|
|
Ok(tokenizer)
|
|
}
|
|
}
|
|
|
|
/// Container for different types of trained tokenizers
|
|
#[derive(Debug, Serialize, Deserialize)]
|
|
pub enum TrainedTokenizer {
|
|
/// Trained BPE tokenizer
|
|
Bpe(BpeTokenizer),
|
|
/// Trained `SentencePiece` tokenizer
|
|
SentencePiece(SentencePieceTokenizer),
|
|
/// Trained multimodal tokenizer
|
|
Multimodal(MultimodalTokenizer),
|
|
}
|
|
|
|
impl TrainedTokenizer {
|
|
/// Get vocabulary size
|
|
#[must_use]
|
|
pub fn vocab_size(&self) -> usize {
|
|
match self {
|
|
Self::Bpe(tokenizer) => tokenizer.vocab_size(),
|
|
Self::SentencePiece(tokenizer) => tokenizer.vocab_size(),
|
|
Self::Multimodal(tokenizer) => tokenizer.vocab_size(),
|
|
}
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn training_config_default() {
|
|
let config = TrainingConfig::default();
|
|
assert_eq!(config.max_iterations, 1000);
|
|
assert_eq!(config.patience, 10);
|
|
assert_eq!(config.validation_split, 0.1);
|
|
assert!(config.verbose);
|
|
assert!(!config.save_checkpoints);
|
|
}
|
|
|
|
#[test]
|
|
fn training_data_creation() {
|
|
let text_data = TrainingData::Text("hello world".to_string());
|
|
match text_data {
|
|
TrainingData::Text(text) => assert_eq!(text, "hello world"),
|
|
_ => panic!("Expected text data"),
|
|
}
|
|
|
|
let file_data = TrainingData::TextFile("test.txt".to_string());
|
|
match file_data {
|
|
TrainingData::TextFile(path) => assert_eq!(path, "test.txt"),
|
|
_ => panic!("Expected file data"),
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn training_algorithm_types() {
|
|
let bpe_alg = TrainingAlgorithm::Bpe(BpeConfig::default());
|
|
match bpe_alg {
|
|
TrainingAlgorithm::Bpe(_) => (),
|
|
_ => panic!("Expected BPE algorithm"),
|
|
}
|
|
|
|
let sp_config = SentencePieceConfig {
|
|
vocab_size: 8000,
|
|
model_type: SentencePieceModelType::Unigram,
|
|
special_tokens: vec!["[UNK]".to_string()],
|
|
character_coverage_scaled: 9995, // 0.9995 * 10000
|
|
input_sentence_size: 10000,
|
|
normalize_text: true,
|
|
add_dummy_prefix: true,
|
|
min_frequency: 2,
|
|
max_subword_length: 16,
|
|
unk_surface: "▁".to_string(),
|
|
};
|
|
let sp_alg = TrainingAlgorithm::SentencePieceUnigram(sp_config);
|
|
match sp_alg {
|
|
TrainingAlgorithm::SentencePieceUnigram(_) => (),
|
|
_ => panic!("Expected SentencePiece Unigram algorithm"),
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn trainer_creation_and_basic_training() {
|
|
let config = TrainingConfig {
|
|
algorithm: TrainingAlgorithm::Bpe(BpeConfig {
|
|
vocab_size: 100,
|
|
min_frequency: 1,
|
|
..BpeConfig::default()
|
|
}),
|
|
max_iterations: 10,
|
|
verbose: false,
|
|
..TrainingConfig::default()
|
|
};
|
|
|
|
let mut trainer = Trainer::new(config);
|
|
let data = TrainingData::Text("hello world test hello".to_string());
|
|
|
|
let trained_tokenizer = trainer.train(data).await.unwrap();
|
|
assert!(trained_tokenizer.vocab_size() > 0);
|
|
|
|
let metrics = trainer.metrics();
|
|
assert!(metrics.training_time_seconds > 0.0);
|
|
assert!(metrics.converged);
|
|
}
|
|
|
|
#[test]
|
|
fn training_metrics_default() {
|
|
let metrics = TrainingMetrics::default();
|
|
assert_eq!(metrics.loss_history.len(), 0);
|
|
assert_eq!(metrics.validation_loss.len(), 0);
|
|
assert_eq!(metrics.vocab_growth.len(), 0);
|
|
assert_eq!(metrics.training_time_seconds, 0.0);
|
|
assert_eq!(metrics.final_vocab_size, 0);
|
|
assert_eq!(metrics.iterations_completed, 0);
|
|
assert!(!metrics.converged);
|
|
}
|
|
}
|