app/fine_tuning/trainer.py

"""
BART Fine-Tuning Engine with LoRA

This module provides fine-tuning capabilities for the BART zero-shot classifier
using Parameter-Efficient Fine-Tuning (PEFT) with LoRA (Low-Rank Adaptation).
"""

import os
import json
import numpy as np
from datetime import datetime
from typing import List, Dict, Tuple, Optional
import warnings

import torch
from transformers import (
    AutoTokenizer,
    AutoModelForSequenceClassification,
    Trainer,
    TrainingArguments,
    EarlyStoppingCallback,
    TrainerCallback,
    TrainerState,
    TrainerControl
)
from peft import LoraConfig, get_peft_model, TaskType
from datasets import Dataset
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score, precision_recall_fscore_support, confusion_matrix
import logging

# Suppress expected warnings
warnings.filterwarnings('ignore', message='.*num_labels.*incompatible.*')
warnings.filterwarnings('ignore', message='.*missing keys.*checkpoint.*')

logger = logging.getLogger(__name__)


class ProgressCallback(TrainerCallback):
    """Callback to track training progress and update database"""

    def __init__(self, run_id: int):
        self.run_id = run_id

    def on_epoch_begin(self, args, state: TrainerState, control: TrainerControl, **kwargs):
        """Called at the beginning of an epoch"""
        try:
            from app import create_app, db
            from app.models.models import FineTuningRun

            app = create_app()
            with app.app_context():
                run = FineTuningRun.query.get(self.run_id)
                if run:
                    run.current_epoch = int(state.epoch) if state.epoch else 0
                    run.progress_message = f"Starting epoch {run.current_epoch + 1}/{run.total_epochs}"
                    db.session.commit()
        except Exception as e:
            logger.error(f"Error updating progress on epoch begin: {e}")

    def on_step_end(self, args, state: TrainerState, control: TrainerControl, **kwargs):
        """Called at the end of a training step"""
        try:
            # Update every 5 steps to avoid too many DB writes
            if state.global_step % 5 == 0:
                from app import create_app, db
                from app.models.models import FineTuningRun

                app = create_app()
                with app.app_context():
                    run = FineTuningRun.query.get(self.run_id)
                    if run:
                        run.current_step = state.global_step
                        run.current_epoch = int(state.epoch) if state.epoch else 0

                        # Get current loss if available
                        if state.log_history:
                            last_log = state.log_history[-1]
                            if 'loss' in last_log:
                                run.current_loss = last_log['loss']

                        # Calculate progress percentage
                        if run.total_steps and run.total_steps > 0:
                            progress_pct = (state.global_step / run.total_steps) * 100
                            run.progress_message = f"Epoch {run.current_epoch + 1}/{run.total_epochs} - Step {state.global_step}/{run.total_steps} ({progress_pct:.1f}%)"
                            if run.current_loss:
                                run.progress_message += f" - Loss: {run.current_loss:.4f}"

                        db.session.commit()
        except Exception as e:
            logger.error(f"Error updating progress on step end: {e}")

    def on_log(self, args, state: TrainerState, control: TrainerControl, logs=None, **kwargs):
        """Called when logging occurs"""
        try:
            from app import create_app, db
            from app.models.models import FineTuningRun

            app = create_app()
            with app.app_context():
                run = FineTuningRun.query.get(self.run_id)
                if run and logs:
                    if 'loss' in logs:
                        run.current_loss = logs['loss']
                    db.session.commit()
        except Exception as e:
            logger.error(f"Error updating progress on log: {e}")


class BARTFineTuner:
    """Fine-tune BART model for multi-class classification using LoRA"""

    def __init__(self, base_model_name: str = "facebook/bart-large-mnli"):
        """
        Initialize the fine-tuner.

        Args:
            base_model_name: Hugging Face model ID for the base model
        """
        self.base_model_name = base_model_name
        self.tokenizer = None
        self.model = None
        self.categories = ['Vision', 'Problem', 'Objectives', 'Directives', 'Values', 'Actions']
        self.label2id = {label: idx for idx, label in enumerate(self.categories)}
        self.id2label = {idx: label for idx, label in enumerate(self.categories)}

    def prepare_dataset(
        self,
        training_examples: List[Dict],
        train_split: float = 0.7,
        val_split: float = 0.15,
        test_split: float = 0.15,
        random_state: int = 42
    ) -> Tuple[Dataset, Dataset, Dataset]:
        """
        Prepare training, validation, and test datasets from training examples.

        Args:
            training_examples: List of dicts with 'message' and 'corrected_category'
            train_split: Proportion for training set
            val_split: Proportion for validation set
            test_split: Proportion for test set
            random_state: Random seed for reproducibility

        Returns:
            Tuple of (train_dataset, val_dataset, test_dataset)
        """
        logger.info(f"Preparing dataset from {len(training_examples)} examples")

        # Extract texts and labels
        texts = [ex['message'] for ex in training_examples]
        labels = [self.label2id[ex['corrected_category']] for ex in training_examples]

        # Validate splits
        assert abs(train_split + val_split + test_split - 1.0) < 0.01, "Splits must sum to 1.0"

        num_classes = len(self.categories)
        total_examples = len(texts)

        # Calculate minimum examples needed for stratified split
        # Need at least num_classes examples in each split
        min_test_size = int(total_examples * test_split)
        min_val_size = int(total_examples * val_split)

        # Check if we have enough examples for stratification
        use_stratify = (min_test_size >= num_classes and min_val_size >= num_classes)

        if not use_stratify:
            logger.warning(f"Dataset too small ({total_examples} examples) for stratified split. "
                          f"Using random split instead.")

        # First split: separate test set
        train_val_texts, test_texts, train_val_labels, test_labels = train_test_split(
            texts, labels,
            test_size=test_split,
            random_state=random_state,
            stratify=labels if use_stratify else None
        )

        # Second split: separate train and validation
        val_size_adjusted = val_split / (train_split + val_split)
        train_texts, val_texts, train_labels, val_labels = train_test_split(
            train_val_texts, train_val_labels,
            test_size=val_size_adjusted,
            random_state=random_state,
            stratify=train_val_labels if use_stratify else None
        )

        # Tokenize datasets
        train_dataset = self._create_dataset(train_texts, train_labels)
        val_dataset = self._create_dataset(val_texts, val_labels)
        test_dataset = self._create_dataset(test_texts, test_labels)

        logger.info(f"Dataset prepared: train={len(train_dataset)}, "
                   f"val={len(val_dataset)}, test={len(test_dataset)}")

        return train_dataset, val_dataset, test_dataset

    def _create_dataset(self, texts: List[str], labels: List[int]) -> Dataset:
        """Create a Hugging Face Dataset with tokenized texts"""
        # Load tokenizer if not already loaded
        if self.tokenizer is None:
            self.tokenizer = AutoTokenizer.from_pretrained(self.base_model_name)

        # Tokenize
        encodings = self.tokenizer(
            texts,
            truncation=True,
            padding='max_length',
            max_length=128,
            return_tensors='pt'
        )

        # Create dataset
        dataset_dict = {
            'input_ids': encodings['input_ids'],
            'attention_mask': encodings['attention_mask'],
            'labels': torch.tensor(labels)
        }

        return Dataset.from_dict(dataset_dict)

    def setup_head_only_model(self) -> None:
        """
        Set up BART model for classification head-only fine-tuning.
        Freezes the encoder and only trains the classification head.
        Better for small datasets (<100 examples).
        """
        logger.info("Setting up BART model for head-only training")

        # Load base model
        self.model = AutoModelForSequenceClassification.from_pretrained(
            self.base_model_name,
            num_labels=len(self.categories),
            id2label=self.id2label,
            label2id=self.label2id,
            problem_type="single_label_classification",
            ignore_mismatched_sizes=True
        )

        # Freeze all parameters except classification head
        for name, param in self.model.named_parameters():
            if 'classification_head' in name or 'classifier' in name:
                param.requires_grad = True
            else:
                param.requires_grad = False

        # Count trainable parameters
        trainable = sum(p.numel() for p in self.model.parameters() if p.requires_grad)
        total = sum(p.numel() for p in self.model.parameters())
        logger.info(f"Trainable params: {trainable:,} / {total:,} ({100 * trainable / total:.2f}%)")

    def setup_lora_model(self, lora_config: Dict) -> None:
        """
        Set up BART model with LoRA adapters.

        Args:
            lora_config: Dict with LoRA hyperparameters:
                - r: Rank of update matrices (default: 16)
                - lora_alpha: Scaling factor (default: 32)
                - lora_dropout: Dropout probability (default: 0.1)
                - target_modules: Modules to apply LoRA to
        """
        logger.info("Setting up BART model with LoRA")

        # Load base model for sequence classification
        self.model = AutoModelForSequenceClassification.from_pretrained(
            self.base_model_name,
            num_labels=len(self.categories),
            id2label=self.id2label,
            label2id=self.label2id,
            problem_type="single_label_classification",
            ignore_mismatched_sizes=True  # BART-MNLI has 3 classes, we need 6
        )

        # Configure LoRA
        peft_config = LoraConfig(
            task_type=TaskType.SEQ_CLS,
            inference_mode=False,
            r=lora_config.get('r', 16),
            lora_alpha=lora_config.get('lora_alpha', 32),
            lora_dropout=lora_config.get('lora_dropout', 0.1),
            target_modules=lora_config.get('target_modules', ['q_proj', 'v_proj']),
            bias="none"
        )

        # Apply PEFT
        self.model = get_peft_model(self.model, peft_config)
        self.model.print_trainable_parameters()

        logger.info("LoRA model ready")

    def train(
        self,
        train_dataset: Dataset,
        val_dataset: Dataset,
        output_dir: str,
        training_config: Dict,
        run_id: Optional[int] = None
    ) -> Dict:
        """
        Train the model with LoRA.

        Args:
            train_dataset: Training dataset
            val_dataset: Validation dataset
            output_dir: Directory to save model checkpoints
            training_config: Training hyperparameters:
                - learning_rate: Learning rate (default: 3e-4)
                - num_epochs: Number of training epochs (default: 3)
                - batch_size: Per-device batch size (default: 8)
                - warmup_ratio: Warmup ratio (default: 0.1)

        Returns:
            Dict with training metrics
        """
        logger.info("Starting training")

        # Create output directory
        os.makedirs(output_dir, exist_ok=True)

        # Force CPU training to avoid cuDNN compatibility issues on WSL2
        use_cuda = False
        logger.info("Using CPU for training (CUDA disabled to avoid compatibility issues)")

        # Training arguments
        training_args = TrainingArguments(
            output_dir=output_dir,
            num_train_epochs=training_config.get('num_epochs', 3),
            per_device_train_batch_size=training_config.get('batch_size', 8),
            per_device_eval_batch_size=training_config.get('batch_size', 8),
            learning_rate=training_config.get('learning_rate', 3e-4),
            warmup_ratio=training_config.get('warmup_ratio', 0.1),
            weight_decay=0.01,
            logging_dir=f'{output_dir}/logs',
            logging_steps=10,
            eval_strategy="epoch",
            save_strategy="epoch",
            load_best_model_at_end=True,
            metric_for_best_model="eval_loss",
            greater_is_better=False,
            save_total_limit=2,
            report_to="none",  # Disable wandb, tensorboard
            use_cpu=not use_cuda,  # Use CPU if CUDA test fails
            fp16=use_cuda,  # Only use mixed precision with working CUDA
        )

        # Calculate total steps for progress tracking
        num_epochs = training_config.get('num_epochs', 3)
        batch_size = training_config.get('batch_size', 8)
        total_steps = (len(train_dataset) // batch_size) * num_epochs

        # Update run with total steps and epochs if run_id provided
        if run_id:
            try:
                from app import create_app, db
                from app.models.models import FineTuningRun

                app = create_app()
                with app.app_context():
                    run = FineTuningRun.query.get(run_id)
                    if run:
                        run.total_epochs = num_epochs
                        run.total_steps = total_steps
                        db.session.commit()
            except Exception as e:
                logger.error(f"Error updating run totals: {e}")

        # Prepare callbacks
        callbacks = [EarlyStoppingCallback(early_stopping_patience=2)]
        if run_id:
            callbacks.append(ProgressCallback(run_id))

        # Trainer
        trainer = Trainer(
            model=self.model,
            args=training_args,
            train_dataset=train_dataset,
            eval_dataset=val_dataset,
            tokenizer=self.tokenizer,
            callbacks=callbacks
        )

        # Train
        train_result = trainer.train()

        # Save model
        trainer.save_model(output_dir)
        self.tokenizer.save_pretrained(output_dir)

        # Extract metrics
        metrics = {
            'train_loss': train_result.metrics.get('train_loss'),
            'train_runtime': train_result.metrics.get('train_runtime'),
            'train_samples_per_second': train_result.metrics.get('train_samples_per_second'),
        }

        # Validation metrics
        eval_metrics = trainer.evaluate()
        metrics['val_loss'] = eval_metrics.get('eval_loss')

        logger.info(f"Training complete: {metrics}")

        return metrics

    def evaluate(
        self,
        test_dataset: Dataset,
        model_path: Optional[str] = None
    ) -> Dict:
        """
        Evaluate model on test set.

        Args:
            test_dataset: Test dataset
            model_path: Path to saved model (if None, uses current model)

        Returns:
            Dict with evaluation metrics
        """
        logger.info("Evaluating model")

        # Load model if path provided
        if model_path and os.path.exists(model_path):
            self.tokenizer = AutoTokenizer.from_pretrained(model_path)
            self.model = AutoModelForSequenceClassification.from_pretrained(
                model_path,
                num_labels=len(self.categories),
                ignore_mismatched_sizes=True
            )

        # Make predictions
        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        self.model.to(device)
        self.model.eval()

        predictions = []
        true_labels = []

        with torch.no_grad():
            for i in range(len(test_dataset)):
                # Get the data - handle both tensor and list formats
                item = test_dataset[i]

                # Convert to tensors if needed
                input_ids = torch.tensor(item['input_ids']) if isinstance(item['input_ids'], list) else item['input_ids']
                attention_mask = torch.tensor(item['attention_mask']) if isinstance(item['attention_mask'], list) else item['attention_mask']
                label = torch.tensor(item['labels']) if isinstance(item['labels'], list) else item['labels']

                # Create batch
                batch = {
                    'input_ids': input_ids.unsqueeze(0).to(device),
                    'attention_mask': attention_mask.unsqueeze(0).to(device)
                }

                outputs = self.model(**batch)
                pred = torch.argmax(outputs.logits, dim=1).item()
                predictions.append(pred)
                true_labels.append(label.item() if isinstance(label, torch.Tensor) else label)

        # Calculate metrics
        accuracy = accuracy_score(true_labels, predictions)
        precision, recall, f1, _ = precision_recall_fscore_support(
            true_labels, predictions, average='macro', zero_division=0
        )

        # Per-category metrics
        precision_per_cat, recall_per_cat, f1_per_cat, _ = precision_recall_fscore_support(
            true_labels, predictions, average=None, zero_division=0, labels=range(len(self.categories))
        )

        per_category_metrics = {}
        for idx, category in enumerate(self.categories):
            per_category_metrics[category] = {
                'precision': float(precision_per_cat[idx]),
                'recall': float(recall_per_cat[idx]),
                'f1': float(f1_per_cat[idx])
            }

        # Confusion matrix
        cm = confusion_matrix(true_labels, predictions, labels=range(len(self.categories)))

        metrics = {
            'test_accuracy': float(accuracy),
            'test_precision_macro': float(precision),
            'test_recall_macro': float(recall),
            'test_f1_macro': float(f1),
            'per_category': per_category_metrics,
            'confusion_matrix': cm.tolist()
        }

        logger.info(f"Evaluation complete: accuracy={accuracy:.3f}, f1={f1:.3f}")

        return metrics

    def compare_to_baseline(
        self,
        test_texts: List[str],
        test_labels: List[str]
    ) -> float:
        """
        Compare fine-tuned model performance to baseline zero-shot classifier.

        Args:
            test_texts: Test text samples
            test_labels: True category labels

        Returns:
            Improvement in accuracy over baseline
        """
        logger.info("Comparing to baseline model")

        # Load baseline zero-shot classifier
        from transformers import pipeline
        baseline_classifier = pipeline(
            "zero-shot-classification",
            model=self.base_model_name,
            device=0 if torch.cuda.is_available() else -1
        )

        # Get baseline predictions
        candidate_labels = [
            f"{cat}: {desc}"
            for cat, desc in zip(
                self.categories,
                [
                    "future aspirations, desired outcomes, what success looks like",
                    "current issues, frustrations, causes of problems",
                    "specific goals to achieve",
                    "restrictions or requirements for solution design",
                    "principles or restrictions for setting objectives",
                    "concrete steps, interventions, or activities to implement"
                ]
            )
        ]

        baseline_preds = []
        for text in test_texts:
            result = baseline_classifier(text, candidate_labels, multi_label=False)
            top_label = result['labels'][0].split(':')[0]
            baseline_preds.append(top_label)

        baseline_accuracy = accuracy_score(test_labels, baseline_preds)

        # Get fine-tuned model predictions (already evaluated)
        # This is a simplified comparison - in practice, reuse evaluation results
        logger.info(f"Baseline accuracy: {baseline_accuracy:.3f}")

        return baseline_accuracy

    def save_metrics(self, metrics: Dict, output_path: str) -> None:
        """Save metrics to JSON file"""
        with open(output_path, 'w') as f:
            json.dump(metrics, f, indent=2)
        logger.info(f"Metrics saved to {output_path}")

    def export_model(self, model_path: str, export_path: str) -> None:
        """
        Export model for deployment or backup.

        Args:
            model_path: Path to saved model
            export_path: Path to export directory
        """
        import shutil

        logger.info(f"Exporting model from {model_path} to {export_path}")

        os.makedirs(export_path, exist_ok=True)

        # Copy model files
        for file in os.listdir(model_path):
            src = os.path.join(model_path, file)
            dst = os.path.join(export_path, file)
            if os.path.isfile(src):
                shutil.copy2(src, dst)

        logger.info("Model exported successfully")