Initial commit
This commit is contained in:
redclawsystems
@@ -0,0 +1,242 @@
|
||||
//! Complete tokenization pipeline
|
||||
//!
|
||||
//! This module provides a unified tokenization pipeline that can combine
|
||||
//! multiple tokenizers and preprocessing steps.
|
||||
|
||||
use crate::{Result, TokenId, TokenizationError, TokenizationStats, Tokenizer};
|
||||
use parking_lot::RwLock;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::sync::Arc;
|
||||
|
||||
/// Pipeline configuration
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct PipelineConfig {
|
||||
/// Name of the pipeline
|
||||
pub name: String,
|
||||
/// Version of the pipeline
|
||||
pub version: String,
|
||||
/// Maximum input length
|
||||
pub max_input_length: usize,
|
||||
/// Whether to truncate long inputs
|
||||
pub truncate: bool,
|
||||
/// Padding token ID
|
||||
pub pad_token_id: Option<TokenId>,
|
||||
}
|
||||
|
||||
/// Preprocessing step
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub enum PreprocessingStep {
|
||||
/// Lowercase the text
|
||||
Lowercase,
|
||||
/// Normalize unicode
|
||||
NormalizeUnicode,
|
||||
/// Remove extra whitespace
|
||||
NormalizeWhitespace,
|
||||
/// Apply custom regex replacement
|
||||
RegexReplace {
|
||||
/// Regular expression pattern to match
|
||||
pattern: String,
|
||||
/// Replacement string
|
||||
replacement: String,
|
||||
},
|
||||
}
|
||||
|
||||
/// Tokenization pipeline
|
||||
#[derive(Debug)]
|
||||
pub struct TokenizationPipeline {
|
||||
config: PipelineConfig,
|
||||
tokenizer: Arc<dyn Tokenizer>,
|
||||
preprocessing_steps: Vec<PreprocessingStep>,
|
||||
stats: Arc<RwLock<TokenizationStats>>,
|
||||
}
|
||||
|
||||
impl TokenizationPipeline {
|
||||
/// Create a new tokenization pipeline
|
||||
pub fn new(
|
||||
config: PipelineConfig,
|
||||
tokenizer: Arc<dyn Tokenizer>,
|
||||
preprocessing_steps: Vec<PreprocessingStep>,
|
||||
) -> Self {
|
||||
Self {
|
||||
config,
|
||||
tokenizer,
|
||||
preprocessing_steps,
|
||||
stats: Arc::new(RwLock::new(TokenizationStats::default())),
|
||||
}
|
||||
}
|
||||
|
||||
/// Apply preprocessing steps to text
|
||||
pub fn preprocess(&self, text: &str) -> Result<String> {
|
||||
let mut processed = text.to_string();
|
||||
|
||||
for step in &self.preprocessing_steps {
|
||||
processed = self.apply_preprocessing_step(&processed, step)?;
|
||||
}
|
||||
|
||||
Ok(processed)
|
||||
}
|
||||
|
||||
/// Apply a single preprocessing step
|
||||
fn apply_preprocessing_step(&self, text: &str, step: &PreprocessingStep) -> Result<String> {
|
||||
match step {
|
||||
PreprocessingStep::Lowercase => Ok(text.to_lowercase()),
|
||||
PreprocessingStep::NormalizeUnicode => {
|
||||
use unicode_normalization::UnicodeNormalization;
|
||||
Ok(text.nfc().collect())
|
||||
}
|
||||
PreprocessingStep::NormalizeWhitespace => {
|
||||
Ok(text.split_whitespace().collect::<Vec<_>>().join(" "))
|
||||
}
|
||||
PreprocessingStep::RegexReplace {
|
||||
pattern,
|
||||
replacement,
|
||||
} => {
|
||||
let regex = regex::Regex::new(pattern)
|
||||
.map_err(|e| TokenizationError::InvalidInput(format!("Invalid regex: {e}")))?;
|
||||
Ok(regex.replace_all(text, replacement).to_string())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Process text through the full pipeline
|
||||
pub async fn process(&self, text: &str) -> Result<Vec<TokenId>> {
|
||||
let start_time = std::time::Instant::now();
|
||||
|
||||
// Apply preprocessing
|
||||
let preprocessed = self.preprocess(text)?;
|
||||
|
||||
// Check length limits
|
||||
if preprocessed.len() > self.config.max_input_length {
|
||||
if self.config.truncate {
|
||||
let truncated = preprocessed
|
||||
.chars()
|
||||
.take(self.config.max_input_length)
|
||||
.collect::<String>();
|
||||
return self.tokenizer.encode(&truncated).await;
|
||||
}
|
||||
return Err(TokenizationError::InvalidInput(format!(
|
||||
"Input length {} exceeds maximum {}",
|
||||
preprocessed.len(),
|
||||
self.config.max_input_length
|
||||
)));
|
||||
}
|
||||
|
||||
// Tokenize
|
||||
let mut token_ids = self.tokenizer.encode(&preprocessed).await?;
|
||||
|
||||
// Apply padding if configured
|
||||
if let Some(pad_token_id) = self.config.pad_token_id {
|
||||
// This is a simple example - real padding would be more sophisticated
|
||||
if token_ids.is_empty() {
|
||||
token_ids.push(pad_token_id);
|
||||
}
|
||||
}
|
||||
|
||||
// Update statistics
|
||||
{
|
||||
let mut stats = self.stats.write();
|
||||
stats.token_count = token_ids.len();
|
||||
stats.processing_time_ms = start_time.elapsed().as_millis() as u64;
|
||||
}
|
||||
|
||||
Ok(token_ids)
|
||||
}
|
||||
|
||||
/// Decode tokens back to text
|
||||
pub async fn decode_tokens(&self, token_ids: &[TokenId]) -> Result<String> {
|
||||
self.tokenizer.decode(token_ids).await
|
||||
}
|
||||
|
||||
/// Get pipeline configuration
|
||||
#[must_use]
|
||||
pub fn get_config(&self) -> &PipelineConfig {
|
||||
&self.config
|
||||
}
|
||||
|
||||
/// Get pipeline statistics
|
||||
#[must_use]
|
||||
pub fn get_stats(&self) -> TokenizationStats {
|
||||
let mut combined_stats = self.stats.read().clone();
|
||||
let tokenizer_stats = self.tokenizer.get_stats();
|
||||
|
||||
// Combine statistics
|
||||
combined_stats.unique_tokens = tokenizer_stats.unique_tokens;
|
||||
combined_stats.avg_token_length = tokenizer_stats.avg_token_length;
|
||||
|
||||
combined_stats
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::bpe::{BpeConfig, BpeTokenizer};
|
||||
|
||||
#[tokio::test]
|
||||
async fn pipeline_creation() {
|
||||
let bpe_config = BpeConfig::default();
|
||||
let tokenizer = Arc::new(BpeTokenizer::new(bpe_config).unwrap());
|
||||
|
||||
let config = PipelineConfig {
|
||||
name: "test_pipeline".to_string(),
|
||||
version: "1.0.0".to_string(),
|
||||
max_input_length: 1000,
|
||||
truncate: true,
|
||||
pad_token_id: Some(1),
|
||||
};
|
||||
|
||||
let steps = vec![
|
||||
PreprocessingStep::Lowercase,
|
||||
PreprocessingStep::NormalizeWhitespace,
|
||||
];
|
||||
|
||||
let pipeline = TokenizationPipeline::new(config, tokenizer, steps);
|
||||
assert_eq!(pipeline.get_config().name, "test_pipeline");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn pipeline_preprocessing() {
|
||||
let bpe_config = BpeConfig::default();
|
||||
let tokenizer = Arc::new(BpeTokenizer::new(bpe_config).unwrap());
|
||||
|
||||
let config = PipelineConfig {
|
||||
name: "test_pipeline".to_string(),
|
||||
version: "1.0.0".to_string(),
|
||||
max_input_length: 1000,
|
||||
truncate: true,
|
||||
pad_token_id: None,
|
||||
};
|
||||
|
||||
let steps = vec![
|
||||
PreprocessingStep::Lowercase,
|
||||
PreprocessingStep::NormalizeWhitespace,
|
||||
];
|
||||
|
||||
let pipeline = TokenizationPipeline::new(config, tokenizer, steps);
|
||||
|
||||
let processed = pipeline.preprocess(" HELLO WORLD ").unwrap();
|
||||
assert_eq!(processed, "hello world");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn preprocessing_step_serialization() {
|
||||
let step = PreprocessingStep::RegexReplace {
|
||||
pattern: r"\d+".to_string(),
|
||||
replacement: "[NUM]".to_string(),
|
||||
};
|
||||
|
||||
let serialized = serde_json::to_string(&step).unwrap();
|
||||
let deserialized: PreprocessingStep = serde_json::from_str(&serialized).unwrap();
|
||||
|
||||
match deserialized {
|
||||
PreprocessingStep::RegexReplace {
|
||||
pattern,
|
||||
replacement,
|
||||
} => {
|
||||
assert_eq!(pattern, r"\d+");
|
||||
assert_eq!(replacement, "[NUM]");
|
||||
}
|
||||
_ => panic!("Wrong preprocessing step type"),
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user