""" Preprocessing Pipeline Runner ================================ Central pipeline that runs all enabled preprocessing stages sequentially and logs each step. """ from dataclasses import dataclass, field from typing import List, Dict, Any, Optional, Tuple import time import pandas as pd from preprocessing.text_cleaning import TextCleaningConfig, apply_text_cleaning from preprocessing.tokenization import ( TokenizationConfig, get_tokenizer, compute_token_stats, truncate_samples, split_long_samples, ) from preprocessing.system_prompt import SystemPromptConfig from preprocessing.dataset_balancing import BalancingConfig, balance_dataset from preprocessing.quality_filters import QualityFilterConfig, apply_quality_filters from preprocessing.deduplication import DeduplicationConfig, apply_deduplication from preprocessing.train_val_split import SplitConfig, split_dataset from preprocessing.output_formatter import OutputFormatConfig, format_dataset, export_jsonl from preprocessing.pii_filter import PIIFilterConfig, apply_pii_filter_df from preprocessing.augmentation import AugmentationConfig, augment_dataset @dataclass class PreprocessingConfig: """Master configuration for the entire preprocessing pipeline.""" # Column mappings instruction_col: str = "" output_col: str = "" input_col: Optional[str] = None label_col: Optional[str] = None # Sub-configs text_cleaning: TextCleaningConfig = field(default_factory=TextCleaningConfig) tokenization: TokenizationConfig = field(default_factory=TokenizationConfig) system_prompt: SystemPromptConfig = field(default_factory=SystemPromptConfig) balancing: BalancingConfig = field(default_factory=BalancingConfig) quality_filters: QualityFilterConfig = field(default_factory=QualityFilterConfig) deduplication: DeduplicationConfig = field(default_factory=DeduplicationConfig) split: SplitConfig = field(default_factory=SplitConfig) output_format: OutputFormatConfig = field(default_factory=OutputFormatConfig) pii_filter: PIIFilterConfig = field(default_factory=PIIFilterConfig) augmentation: AugmentationConfig = field(default_factory=AugmentationConfig) @dataclass class PipelineLog: """A single log entry from a pipeline stage.""" stage: str description: str rows_before: int rows_after: int duration_ms: float @property def rows_delta(self) -> int: return self.rows_after - self.rows_before class PreprocessingPipeline: """ Sequential preprocessing pipeline runner. Applies all enabled stages and collects logs. """ def __init__(self, config: PreprocessingConfig): self.config = config self.logs: List[PipelineLog] = [] def _log(self, stage: str, desc: str, before: int, after: int, elapsed: float): self.logs.append(PipelineLog( stage=stage, description=desc, rows_before=before, rows_after=after, duration_ms=round(elapsed * 1000, 1), )) def run( self, df: pd.DataFrame, progress_callback=None, ) -> Tuple[pd.DataFrame, pd.DataFrame, List[PipelineLog]]: """ Run the complete preprocessing pipeline. Args: df: Input DataFrame progress_callback: Optional callable(stage_name, progress_pct) for UI updates Returns: (train_df, val_df, logs) If split is disabled, val_df will be empty. """ self.logs = [] total_stages = 7 # text cleaning, quality, dedup, pii, balancing, augmentation, tokenization current_stage = 0 def _progress(name): nonlocal current_stage current_stage += 1 if progress_callback: pct = int((current_stage / total_stages) * 100) progress_callback(name, pct) cfg = self.config text_cols = [c for c in [cfg.instruction_col, cfg.output_col, cfg.input_col] if c and c in df.columns] # ── Stage 1: Text Cleaning ── t0 = time.time() before = len(df) any_cleaning = ( cfg.text_cleaning.remove_html or cfg.text_cleaning.remove_urls or cfg.text_cleaning.remove_emojis or cfg.text_cleaning.normalize_whitespace or cfg.text_cleaning.lowercase or cfg.text_cleaning.remove_special_chars or cfg.text_cleaning.strip_extra_linebreaks ) if any_cleaning: df = apply_text_cleaning(df, text_cols, cfg.text_cleaning) self._log("Text Cleaning", "Applied text cleaning filters", before, len(df), time.time() - t0) _progress("Text Cleaning") # ── Stage 2: Quality Filters ── t0 = time.time() before = len(df) has_quality = ( cfg.quality_filters.min_word_count > 0 or cfg.quality_filters.max_word_count > 0 or cfg.quality_filters.profanity_filter or cfg.quality_filters.language_filter or cfg.quality_filters.remove_low_quality ) if has_quality and cfg.output_col: df = apply_quality_filters(df, cfg.output_col, cfg.quality_filters) self._log("Quality Filters", "Applied quality filters", before, len(df), time.time() - t0) _progress("Quality Filters") # ── Stage 3: Deduplication ── t0 = time.time() before = len(df) if cfg.instruction_col and (cfg.deduplication.remove_exact or cfg.deduplication.remove_semantic): df = apply_deduplication(df, cfg.instruction_col, cfg.deduplication) self._log("Deduplication", "Removed duplicate samples", before, len(df), time.time() - t0) _progress("Deduplication") # ── Stage 4: PII Filtering ── t0 = time.time() before = len(df) has_pii = ( cfg.pii_filter.filter_emails or cfg.pii_filter.filter_phones or cfg.pii_filter.filter_id_numbers or cfg.pii_filter.filter_api_keys or cfg.pii_filter.filter_addresses ) if has_pii: df = apply_pii_filter_df(df, text_cols, cfg.pii_filter) self._log("PII Filtering", "Masked PII data", before, len(df), time.time() - t0) _progress("PII Filtering") # ── Stage 5: Dataset Balancing ── t0 = time.time() before = len(df) if cfg.balancing.enabled and cfg.balancing.label_column and cfg.balancing.strategy != "none": df = balance_dataset(df, cfg.balancing.label_column, cfg.balancing.strategy) self._log("Balancing", "Balanced dataset classes", before, len(df), time.time() - t0) _progress("Balancing") # ── Stage 6: Augmentation ── t0 = time.time() before = len(df) if cfg.augmentation.enabled and cfg.instruction_col: df = augment_dataset(df, cfg.instruction_col, cfg.augmentation) self._log("Augmentation", "Generated augmented samples", before, len(df), time.time() - t0) _progress("Augmentation") # ── Stage 7: Tokenization Controls ── t0 = time.time() before = len(df) if cfg.tokenization.truncate_long or cfg.tokenization.split_long: try: tokenizer = get_tokenizer(cfg.tokenization) is_tiktoken = cfg.tokenization.tokenizer_name == "tiktoken" for col in text_cols: if cfg.tokenization.split_long: df = split_long_samples( df, col, cfg.tokenization.max_total_tokens, tokenizer, is_tiktoken, cfg.tokenization.split_overlap, ) elif cfg.tokenization.truncate_long: df = truncate_samples( df, col, cfg.tokenization.max_total_tokens, tokenizer, is_tiktoken, ) except ImportError: pass # tokenizer not available self._log("Tokenization", "Applied tokenization controls", before, len(df), time.time() - t0) _progress("Tokenization") # ── Split ── train_df, val_df = split_dataset(df, cfg.split) return train_df, val_df, self.logs def get_safe_preset() -> PreprocessingConfig: """Return a sensible 'safe preset' configuration for common use cases.""" return PreprocessingConfig( text_cleaning=TextCleaningConfig( remove_html=True, remove_urls=True, remove_emojis=False, normalize_whitespace=True, lowercase=False, remove_special_chars=False, strip_extra_linebreaks=True, ), quality_filters=QualityFilterConfig( min_word_count=3, max_word_count=0, profanity_filter=False, language_filter=False, remove_low_quality=True, min_quality_length=20, ), deduplication=DeduplicationConfig( remove_exact=True, remove_semantic=False, ), pii_filter=PIIFilterConfig( filter_emails=True, filter_phones=True, filter_id_numbers=True, filter_api_keys=True, filter_addresses=False, ), split=SplitConfig( enabled=True, train_ratio=0.9, random_seed=42, shuffle=True, ), output_format=OutputFormatConfig( format_type="openai_chat", ), system_prompt=SystemPromptConfig( system_prompt="You are a helpful AI assistant.", prepend_to_all=True, ), )