data/preprocessing.py

import os
import torch
from torch.utils.data import Dataset, DataLoader
from torch.nn.utils.rnn import pad_sequence
import sentencepiece as spm
from typing import List, Tuple, Optional, Dict
import yaml
import numpy as np
from collections import Counter
import logging

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class TranslationDataset(Dataset):
    def __init__(self, data: List[Tuple[str, str]], tokenizer_path: str, 
                 max_length: int = 100, config_path: str = "configs/config.yaml"):
        """
        Translation dataset for German-English pairs
        
        Args:
            data: List of (source, target) text pairs
            tokenizer_path: Path to SentencePiece model
            max_length: Maximum sequence length
            config_path: Path to config file
        """
        self.data = data
        self.max_length = max_length
        
        # Load config
        with open(config_path, 'r') as f:
            self.config = yaml.safe_load(f)
        
        # Load tokenizer
        self.sp = spm.SentencePieceProcessor()
        self.sp.load(tokenizer_path)
        
        # Special tokens
        self.pad_id = self.sp.pad_id()
        self.bos_id = self.sp.bos_id()
        self.eos_id = self.sp.eos_id()
        self.unk_id = self.sp.unk_id()
        
        logger.info(f"Dataset initialized with {len(self.data)} samples")
        logger.info(f"Vocab size: {self.sp.vocab_size()}")
        
    def __len__(self) -> int:
        return len(self.data)
    
    def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
        src_text, tgt_text = self.data[idx]
        
        # Tokenize
        src_tokens = self.sp.encode(src_text, out_type=int)
        tgt_tokens = self.sp.encode(tgt_text, out_type=int)
        
        # Truncate if necessary
        src_tokens = src_tokens[:self.max_length - 2]  # Leave room for BOS/EOS
        tgt_tokens = tgt_tokens[:self.max_length - 2]
        
        # Add BOS and EOS tokens
        src_tokens = [self.bos_id] + src_tokens + [self.eos_id]
        tgt_tokens = [self.bos_id] + tgt_tokens + [self.eos_id]
        
        # Convert to tensors
        src_tensor = torch.tensor(src_tokens, dtype=torch.long)
        tgt_tensor = torch.tensor(tgt_tokens, dtype=torch.long)
        
        return {
            'src': src_tensor,
            'tgt': tgt_tensor,
            'src_len': len(src_tokens),
            'tgt_len': len(tgt_tokens)
        }


class DataCollator:
    def __init__(self, pad_id: int = 0):
        """
        Collator for batching translation data
        
        Args:
            pad_id: Padding token ID
        """
        self.pad_id = pad_id
    
    def __call__(self, batch: List[Dict[str, torch.Tensor]]) -> Dict[str, torch.Tensor]:
        # Extract sequences
        src_seqs = [item['src'] for item in batch]
        tgt_seqs = [item['tgt'] for item in batch]
        
        # Pad sequences
        src_padded = pad_sequence(src_seqs, batch_first=True, padding_value=self.pad_id)
        tgt_padded = pad_sequence(tgt_seqs, batch_first=True, padding_value=self.pad_id)
        
        # Create attention masks (1 for real tokens, 0 for padding)
        src_mask = (src_padded != self.pad_id).float()
        tgt_mask = (tgt_padded != self.pad_id).float()
        
        return {
            'src': src_padded,
            'tgt': tgt_padded,
            'src_mask': src_mask,
            'tgt_mask': tgt_mask
        }


def create_dataloaders(train_data: List[Tuple[str, str]], 
                      valid_data: List[Tuple[str, str]], 
                      test_data: List[Tuple[str, str]],
                      tokenizer_path: str,
                      batch_size: int = 32,
                      num_workers: int = 2,
                      config_path: str = "configs/config.yaml") -> Tuple[DataLoader, DataLoader, DataLoader]:
    """
    Create DataLoaders for train, validation, and test sets
    
    Args:
        train_data: Training data
        valid_data: Validation data
        test_data: Test data
        tokenizer_path: Path to tokenizer model
        batch_size: Batch size
        num_workers: Number of workers for DataLoader
        config_path: Path to config file
    
    Returns:
        Tuple of (train_loader, valid_loader, test_loader)
    """
    # Create datasets
    train_dataset = TranslationDataset(train_data, tokenizer_path, config_path=config_path)
    valid_dataset = TranslationDataset(valid_data, tokenizer_path, config_path=config_path)
    test_dataset = TranslationDataset(test_data, tokenizer_path, config_path=config_path)
    
    # Create collator
    collator = DataCollator(pad_id=train_dataset.pad_id)
    
    # Create dataloaders
    train_loader = DataLoader(
        train_dataset,
        batch_size=batch_size,
        shuffle=True,
        collate_fn=collator,
        num_workers=num_workers,
        pin_memory=True
    )
    
    valid_loader = DataLoader(
        valid_dataset,
        batch_size=batch_size,
        shuffle=False,
        collate_fn=collator,
        num_workers=num_workers,
        pin_memory=True
    )
    
    test_loader = DataLoader(
        test_dataset,
        batch_size=batch_size,
        shuffle=False,
        collate_fn=collator,
        num_workers=num_workers,
        pin_memory=True
    )
    
    return train_loader, valid_loader, test_loader


def analyze_dataset(data: List[Tuple[str, str]], tokenizer_path: str) -> Dict:
    """
    Analyze dataset statistics
    
    Args:
        data: List of (source, target) pairs
        tokenizer_path: Path to tokenizer
    
    Returns:
        Dictionary with statistics
    """
    sp = spm.SentencePieceProcessor()
    sp.load(tokenizer_path)
    
    src_lengths = []
    tgt_lengths = []
    
    for src, tgt in data:
        src_tokens = sp.encode(src)
        tgt_tokens = sp.encode(tgt)
        src_lengths.append(len(src_tokens))
        tgt_lengths.append(len(tgt_tokens))
    
    stats = {
        'num_samples': len(data),
        'src_avg_length': np.mean(src_lengths),
        'src_max_length': np.max(src_lengths),
        'src_min_length': np.min(src_lengths),
        'tgt_avg_length': np.mean(tgt_lengths),
        'tgt_max_length': np.max(tgt_lengths),
        'tgt_min_length': np.min(tgt_lengths),
        'vocab_size': sp.vocab_size()
    }
    
    return stats


if __name__ == "__main__":
    # Test the dataset
    from data.download import DataDownloader
    
    downloader = DataDownloader()
    train_data, valid_data, test_data = downloader.download_multi30k()
    
    if train_data:
        tokenizer_path = os.path.join('data', 'processed', 'tokenizer.model')
        
        # Analyze dataset
        stats = analyze_dataset(train_data, tokenizer_path)
        logger.info("Dataset statistics:")
        for key, value in stats.items():
            logger.info(f"{key}: {value}")
        
        # Create dataloaders
        train_loader, valid_loader, test_loader = create_dataloaders(
            train_data[:100],  # Use small subset for testing
            valid_data[:10],
            test_data[:10],
            tokenizer_path,
            batch_size=8
        )
        
        # Test loading a batch
        for batch in train_loader:
            logger.info(f"Batch shapes:")
            logger.info(f"  src: {batch['src'].shape}")
            logger.info(f"  tgt: {batch['tgt'].shape}")
            logger.info(f"  src_mask: {batch['src_mask'].shape}")
            logger.info(f"  tgt_mask: {batch['tgt_mask'].shape}")
            break