data_preprocessing.py

"""
Data preprocessing for ModernBERT bias classification training
"""

import torch
import numpy as np
from datasets import load_dataset
from transformers import AutoTokenizer
from sklearn.model_selection import train_test_split
import config
import pandas as pd
from typing import Dict, List, Tuple, Any


def set_seed(seed: int = config.SEED):
    """Set random seed for reproducibility"""
    torch.manual_seed(seed)
    np.random.seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)


def load_and_split_dataset() -> Tuple[Any, Any]:
    """Load dataset and create train/validation split"""
    print("Loading dataset...")
    dataset = load_dataset(config.DATASET_NAME)

    # Get the train split
    train_data = dataset['train']

    # Convert to pandas for easier manipulation
    df = train_data.to_pandas()

    # Create train/validation split
    train_df, val_df = train_test_split(
        df,
        test_size=config.VALIDATION_SPLIT,
        random_state=config.RANDOM_SEED,
        stratify=df[config.LABEL_NAMES].values.argmax(
            axis=1)  # Stratify by dominant label
    )

    print(f"Training examples: {len(train_df)}")
    print(f"Validation examples: {len(val_df)}")

    return train_df, val_df


def prepare_labels(df: pd.DataFrame) -> torch.Tensor:
    """Extract multi-label targets from dataframe"""
    labels = df[config.LABEL_NAMES].values.astype(float)
    return torch.tensor(labels, dtype=torch.float32)


def tokenize_texts(texts: List[str], tokenizer: AutoTokenizer) -> Dict[str, torch.Tensor]:
    """Tokenize input texts"""
    encoded = tokenizer(
        texts,
        truncation=True,
        padding=True,
        max_length=config.MAX_LENGTH,
        return_tensors="pt"
    )
    return encoded


class BiasDataset(torch.utils.data.Dataset):
    """Custom dataset for multi-label bias classification"""

    def __init__(self, texts: List[str], labels: torch.Tensor, tokenizer: AutoTokenizer):
        self.texts = texts
        self.labels = labels
        self.tokenizer = tokenizer

    def __len__(self):
        return len(self.texts)

    def __getitem__(self, idx):
        text = self.texts[idx]
        label = self.labels[idx]

        # Tokenize the text
        encoded = self.tokenizer(
            text,
            truncation=True,
            padding="max_length",
            max_length=config.MAX_LENGTH,
            return_tensors="pt"
        )

        return {
            'input_ids': encoded['input_ids'].squeeze(0),
            'attention_mask': encoded['attention_mask'].squeeze(0),
            'labels': label
        }


def create_data_loaders(tokenizer: AutoTokenizer) -> Tuple[torch.utils.data.DataLoader, torch.utils.data.DataLoader]:
    """Create training and validation data loaders"""
    # Load and split data
    train_df, val_df = load_and_split_dataset()

    # Prepare texts and labels
    train_texts = train_df[config.TEXT_COLUMN].tolist()
    val_texts = val_df[config.TEXT_COLUMN].tolist()

    train_labels = prepare_labels(train_df)
    val_labels = prepare_labels(val_df)

    # Create datasets
    train_dataset = BiasDataset(train_texts, train_labels, tokenizer)
    val_dataset = BiasDataset(val_texts, val_labels, tokenizer)

    # Create data loaders
    train_loader = torch.utils.data.DataLoader(
        train_dataset,
        batch_size=config.BATCH_SIZE,
        shuffle=True,
        num_workers=config.DATALOADER_NUM_WORKERS,
        pin_memory=config.DATALOADER_PIN_MEMORY
    )

    val_loader = torch.utils.data.DataLoader(
        val_dataset,
        batch_size=config.BATCH_SIZE,
        shuffle=False,
        num_workers=config.DATALOADER_NUM_WORKERS,
        pin_memory=config.DATALOADER_PIN_MEMORY
    )

    return train_loader, val_loader


def analyze_dataset_distribution():
    """Analyze the distribution of labels in the dataset"""
    print("Analyzing dataset distribution...")
    dataset = load_dataset(config.DATASET_NAME)
    df = dataset['train'].to_pandas()

    print("\nLabel distribution:")
    for label in config.LABEL_NAMES:
        positive_count = df[label].sum()
        percentage = (positive_count / len(df)) * 100
        print(f"  {label}: {positive_count} ({percentage:.1f}%)")

    # Multi-label statistics
    label_counts = df[config.LABEL_NAMES].sum(axis=1)
    print(f"\nMulti-label statistics:")
    print(f"  Average labels per example: {label_counts.mean():.2f}")
    print(f"  Max labels per example: {label_counts.max()}")
    print(f"  Examples with no labels: {(label_counts == 0).sum()}")
    print(f"  Examples with multiple labels: {(label_counts > 1).sum()}")


if __name__ == "__main__":
    set_seed()
    analyze_dataset_distribution()

    # Test data loading
    print("\nTesting data loading...")
    tokenizer = AutoTokenizer.from_pretrained(config.MODEL_NAME)
    train_loader, val_loader = create_data_loaders(tokenizer)

    # Test first batch
    batch = next(iter(train_loader))
    print(f"Batch input_ids shape: {batch['input_ids'].shape}")
    print(f"Batch attention_mask shape: {batch['attention_mask'].shape}")
    print(f"Batch labels shape: {batch['labels'].shape}")
    print(f"Sample labels: {batch['labels'][0]}")