rustytorch/crates/core/rtx-tokenization/src/pipeline.rs

//! 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"),
        }
    }
}