train_model.py

"""
Train a DistilBERT model for ad detection using transcripts and ad segments.
Exports model to Core ML format for iPhone deployment.
"""

import json
import os
from pathlib import Path
import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.metrics import classification_report, confusion_matrix, accuracy_score, precision_recall_fscore_support
import torch
from torch.utils.data import Dataset, DataLoader
from transformers import (
    DistilBertForSequenceClassification,
    DistilBertTokenizer,
    Trainer,
    TrainingArguments,
    EarlyStoppingCallback
)
from datasets import Dataset as HFDataset
import config

class AdDataset(Dataset):
    """Dataset class for ad detection training."""
    
    def __init__(self, texts, labels, tokenizer, max_length=512):
        self.texts = texts
        self.labels = labels
        self.tokenizer = tokenizer
        self.max_length = max_length
    
    def __len__(self):
        return len(self.texts)
    
    def __getitem__(self, idx):
        text = str(self.texts[idx])
        label = self.labels[idx]
        
        encoding = self.tokenizer(
            text,
            truncation=True,
            padding='max_length',
            max_length=self.max_length,
            return_tensors='pt'
        )
        
        return {
            'input_ids': encoding['input_ids'].flatten(),
            'attention_mask': encoding['attention_mask'].flatten(),
            'labels': torch.tensor(label, dtype=torch.long)
        }

def load_metadata():
    """Load metadata.json if it exists."""
    if config.METADATA_FILE.exists():
        try:
            with open(config.METADATA_FILE, 'r', encoding='utf-8') as f:
                return json.load(f)
        except Exception as e:
            print(f"Error loading metadata: {e}")
    
    return {"episodes": {}}

def prepare_training_data():
    """Prepare training data from transcripts and ad segments."""
    metadata = load_metadata()
    
    texts = []
    labels = []
    
    print("Preparing training data...")
    
    for episode_key, episode_data in metadata.get('episodes', {}).items():
        if not episode_data.get('ad_detected', False):
            continue
        
        transcript_path = config.BASE_DIR / episode_data.get('transcript_path', '')
        ad_path = config.BASE_DIR / episode_data.get('ad_path', '')
        
        if not transcript_path.exists():
            continue
        
        # Load transcript
        with open(transcript_path, 'r', encoding='utf-8') as f:
            transcript_data = json.load(f)
        
        # Load ad segments
        ad_segments = []
        if ad_path and ad_path.exists():
            with open(ad_path, 'r', encoding='utf-8') as f:
                ad_data = json.load(f)
                ad_segments = ad_data.get('ad_segments', [])
        
        # Create time-based mapping of ad segments
        ad_ranges = []
        for ad in ad_segments:
            ad_ranges.append({
                'start': ad.get('start_time', 0),
                'end': ad.get('end_time', 0),
                'text': ad.get('text', '')
            })
        
        # Extract segments from transcript
        segments = transcript_data.get('segments', [])
        if not segments:
            # If no segments, use full text as one segment
            full_text = transcript_data.get('text', '')
            if full_text:
                segments = [{'start': 0, 'end': 0, 'text': full_text}]
        
        # Label each segment as ad or content
        for segment in segments:
            segment_start = segment.get('start', 0)
            segment_end = segment.get('end', segment_start + 5)  # Default 5 seconds if no end
            segment_text = segment.get('text', '').strip()
            
            if not segment_text or len(segment_text) < 10:
                continue
            
            # Check if this segment overlaps with any ad
            is_ad = False
            for ad_range in ad_ranges:
                # Check for overlap
                if (segment_start < ad_range['end'] and segment_end > ad_range['start']):
                    is_ad = True
                    break
            
            # Use sliding windows for longer segments
            if len(segment_text) > 512:
                # Split into overlapping windows
                words = segment_text.split()
                window_size = 400  # ~512 tokens
                overlap = 100
                
                for i in range(0, len(words), window_size - overlap):
                    window_text = ' '.join(words[i:i+window_size])
                    if len(window_text) > 20:
                        texts.append(window_text)
                        labels.append(1 if is_ad else 0)
            else:
                texts.append(segment_text)
                labels.append(1 if is_ad else 0)
    
    print(f"Total samples: {len(texts)}")
    print(f"Ads: {sum(labels)}, Content: {len(labels) - sum(labels)}")
    
    return texts, labels

def balance_dataset(texts, labels):
    """Balance dataset by downsampling majority class."""
    texts = np.array(texts, dtype=object)
    labels = np.array(labels)
    
    ad_indices = np.where(labels == 1)[0]
    content_indices = np.where(labels == 0)[0]
    
    # Downsample majority class
    min_class_size = min(len(ad_indices), len(content_indices))
    
    if len(ad_indices) > min_class_size:
        np.random.seed(42)
        selected_ad_indices = np.random.choice(ad_indices, min_class_size, replace=False)
    else:
        selected_ad_indices = ad_indices
    
    if len(content_indices) > min_class_size:
        np.random.seed(42)
        selected_content_indices = np.random.choice(content_indices, min_class_size, replace=False)
    else:
        selected_content_indices = content_indices
    
    balanced_indices = np.concatenate([selected_ad_indices, selected_content_indices])
    np.random.shuffle(balanced_indices)
    
    return texts[balanced_indices].tolist(), labels[balanced_indices].tolist()

def train_model():
    """Train DistilBERT model for ad detection."""
    # Prepare data
    texts, labels = prepare_training_data()
    
    if len(texts) == 0:
        print("Error: No training data available. Run download, transcription, and ad detection first.")
        return
    
    # Balance dataset
    texts, labels = balance_dataset(texts, labels)
    
    print(f"\nBalanced dataset:")
    print(f"Total samples: {len(texts)}")
    print(f"Ads: {sum(labels)}, Content: {len(labels) - sum(labels)}")
    
    # Split into train/test
    X_train, X_test, y_train, y_test = train_test_split(
        texts, labels, test_size=0.2, random_state=42, stratify=labels
    )
    
    print(f"\nTrain set: {len(X_train)} samples")
    print(f"Test set: {len(X_test)} samples")
    
    # Load tokenizer and model
    model_name = "distilbert-base-multilingual-cased"
    tokenizer = DistilBertTokenizer.from_pretrained(model_name)
    model = DistilBertForSequenceClassification.from_pretrained(
        model_name,
        num_labels=2
    )
    
    # Create datasets
    train_dataset = HFDataset.from_dict({
        'text': X_train,
        'label': y_train
    })
    
    test_dataset = HFDataset.from_dict({
        'text': X_test,
        'label': y_test
    })
    
    def tokenize_function(examples):
        return tokenizer(
            examples['text'],
            truncation=True,
            padding='max_length',
            max_length=config.MAX_SEQUENCE_LENGTH
        )
    
    train_dataset = train_dataset.map(tokenize_function, batched=True)
    test_dataset = test_dataset.map(tokenize_function, batched=True)
    
    train_dataset.set_format('torch', columns=['input_ids', 'attention_mask', 'label'])
    test_dataset.set_format('torch', columns=['input_ids', 'attention_mask', 'label'])
    
    # Training arguments
    training_args = TrainingArguments(
        output_dir=str(config.MODEL_DIR / "checkpoints"),
        num_train_epochs=config.NUM_EPOCHS,
        per_device_train_batch_size=config.BATCH_SIZE,
        per_device_eval_batch_size=config.BATCH_SIZE,
        learning_rate=config.LEARNING_RATE,
        weight_decay=0.01,
        logging_dir=str(config.MODEL_DIR / "logs"),
        logging_steps=50,
        evaluation_strategy="epoch",
        save_strategy="epoch",
        load_best_model_at_end=True,
        metric_for_best_model="f1",
        greater_is_better=True,
        save_total_limit=2,
        seed=42
    )
    
    # Compute metrics function
    def compute_metrics(eval_pred):
        predictions, labels = eval_pred
        predictions = np.argmax(predictions, axis=1)
        
        precision, recall, f1, _ = precision_recall_fscore_support(
            labels, predictions, average='binary', zero_division=0
        )
        accuracy = accuracy_score(labels, predictions)
        
        return {
            'accuracy': accuracy,
            'precision': precision,
            'recall': recall,
            'f1': f1
        }
    
    # Initialize trainer
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
        eval_dataset=test_dataset,
        compute_metrics=compute_metrics,
        callbacks=[EarlyStoppingCallback(early_stopping_patience=2)]
    )
    
    # Train
    print("\n=== Training Model ===")
    trainer.train()
    
    # Evaluate
    print("\n=== Evaluating Model ===")
    eval_results = trainer.evaluate()
    print(f"Test Results: {eval_results}")
    
    # Final evaluation on test set
    predictions = trainer.predict(test_dataset)
    pred_labels = np.argmax(predictions.predictions, axis=1)
    
    print("\n=== Detailed Classification Report ===")
    print(classification_report(y_test, pred_labels, target_names=['Content', 'Ad']))
    print("\n=== Confusion Matrix ===")
    print(confusion_matrix(y_test, pred_labels))
    
    # Save model
    final_model_path = config.MODEL_DIR / "final_model"
    trainer.save_model(str(final_model_path))
    tokenizer.save_pretrained(str(final_model_path))
    
    print(f"\n✓ Model saved to: {final_model_path}")
    
    # Export to Core ML
    export_to_coreml(final_model_path)

def export_to_coreml(model_path):
    """Export trained model to Core ML format for iPhone deployment."""
    try:
        import coremltools as ct
        from transformers import pipeline
        
        print("\n=== Exporting to Core ML ===")
        
        # Load model and tokenizer
        tokenizer = DistilBertTokenizer.from_pretrained(str(model_path))
        model = DistilBertForSequenceClassification.from_pretrained(str(model_path))
        
        # Create text classification pipeline
        classifier = pipeline(
            "text-classification",
            model=model,
            tokenizer=tokenizer,
            device=0 if torch.cuda.is_available() else -1
        )
        
        # Convert to Core ML
        # Note: This is a simplified conversion. For production, you may need
        # to convert the model architecture more carefully
        coreml_model = ct.convert(
            classifier.model,
            inputs=[
                ct.TensorType(name="input_ids", shape=(1, config.MAX_SEQUENCE_LENGTH), dtype=np.int64),
                ct.TensorType(name="attention_mask", shape=(1, config.MAX_SEQUENCE_LENGTH), dtype=np.int64)
            ],
            outputs=[
                ct.TensorType(name="logits", dtype=np.float32)
            ]
        )
        
        # Add metadata
        coreml_model.author = "CleanCast ML Model"
        coreml_model.short_description = "DistilBERT-based podcast ad detector"
        coreml_model.version = "1.0"
        
        # Save Core ML model
        coreml_path = config.MODEL_DIR / "ad_detector.mlpackage"
        coreml_model.save(str(coreml_path))
        
        print(f"✓ Core ML model saved to: {coreml_path}")
        print("  Note: You may need to fine-tune the Core ML export for your specific use case.")
        
    except Exception as e:
        print(f"Warning: Could not export to Core ML: {e}")
        print("  Model saved in PyTorch format. You can convert manually later.")
        print("  Consider using coremltools or onnx for conversion.")

if __name__ == "__main__":
    train_model()