train_utils.py

# train_utils.py

import torch
import torch.nn as nn
from torch.optim import AdamW
from sklearn.metrics import classification_report
from sklearn.utils.class_weight import compute_class_weight
import numpy as np
from tqdm import tqdm
import pandas as pd
import os
import joblib

from config import DEVICE, LABEL_COLUMNS, NUM_EPOCHS, LEARNING_RATE, MODEL_SAVE_DIR

def get_class_weights(data_df, field, label_encoder):
    """
    Computes balanced class weights for a given target field.
    These weights can be used in the loss function to mitigate class imbalance.
    Args:
        data_df (pd.DataFrame): The DataFrame containing the original (unencoded) label data.
        field (str): The name of the label column for which to compute weights.
        label_encoder (sklearn.preprocessing.LabelEncoder): The label encoder fitted for this field.
    Returns:
        torch.Tensor: A tensor of class weights for the specified field.
    """
    # Get the original labels for the specified field
    y = data_df[field].values
    # Use label_encoder.transform directly - it will handle unseen labels
    try:
        y_encoded = label_encoder.transform(y)
    except ValueError as e:
        print(f"Warning: {e}")
        print(f"Using only seen labels for class weights calculation")
        # Filter out unseen labels
        seen_labels = set(label_encoder.classes_)
        y_filtered = [label for label in y if label in seen_labels]
        y_encoded = label_encoder.transform(y_filtered)
    
    # Ensure y_encoded is integer type
    y_encoded = y_encoded.astype(int)
    
    # Initialize counts for all possible classes
    n_classes = len(label_encoder.classes_)
    class_counts = np.zeros(n_classes, dtype=int)
    
    # Count occurrences of each class
    for i in range(n_classes):
        class_counts[i] = np.sum(y_encoded == i)
    
    # Calculate weights for all classes
    total_samples = len(y_encoded)
    class_weights = np.ones(n_classes)  # Default weight of 1 for unseen classes
    seen_classes = class_counts > 0
    if np.any(seen_classes):
        class_weights[seen_classes] = total_samples / (np.sum(seen_classes) * class_counts[seen_classes])
    
    return torch.tensor(class_weights, dtype=torch.float)

def initialize_criterions(data_df, label_encoders):
    """
    Initializes CrossEntropyLoss criteria for each label column, applying class weights.
    Args:
        data_df (pd.DataFrame): The original (unencoded) DataFrame. Used to compute class weights.
        label_encoders (dict): Dictionary of LabelEncoder objects.
    Returns:
        dict: A dictionary where keys are label column names and values are
              initialized `torch.nn.CrossEntropyLoss` objects.
    """
    field_criterions = {}
    for field in LABEL_COLUMNS:
        # Get class weights for the current field
        weights = get_class_weights(data_df, field, label_encoders[field])
        # Initialize CrossEntropyLoss with the computed weights and move to the device
        field_criterions[field] = torch.nn.CrossEntropyLoss(weight=weights.to(DEVICE))
    return field_criterions

def train_model(model, loader, optimizer, field_criterions, epoch):
    """
    Trains the given PyTorch model for one epoch.
    Args:
        model (torch.nn.Module): The model to train.
        loader (torch.utils.data.DataLoader): DataLoader for training data.
        optimizer (torch.optim.Optimizer): Optimizer for model parameters.
        field_criterions (dict): Dictionary of loss functions for each label.
        epoch (int): Current epoch number (for progress bar description).
    Returns:
        float: Average training loss for the epoch.
    """
    model.train() # Set the model to training mode
    total_loss = 0
    # Use tqdm for a progress bar during training
    tqdm_loader = tqdm(loader, desc=f"Epoch {epoch + 1} Training")

    for batch in tqdm_loader:
        # Unpack batch based on whether it contains metadata
        if len(batch) == 2: # Text-only models (inputs, labels)
            inputs, labels = batch
            input_ids = inputs['input_ids'].to(DEVICE)
            attention_mask = inputs['attention_mask'].to(DEVICE)
            labels = labels.to(DEVICE)
            # Forward pass through the model
            outputs = model(input_ids, attention_mask)
        elif len(batch) == 3: # Text + Metadata models (inputs, metadata, labels)
            inputs, metadata, labels = batch
            input_ids = inputs['input_ids'].to(DEVICE)
            attention_mask = inputs['attention_mask'].to(DEVICE)
            metadata = metadata.to(DEVICE)
            labels = labels.to(DEVICE)
            # Forward pass through the hybrid model
            outputs = model(input_ids, attention_mask, metadata)
        else:
            raise ValueError("Unsupported batch format. Expected 2 or 3 items in batch.")

        loss = 0
        # Calculate total loss by summing loss for each label column
        # `outputs` is a list of logits, one for each label column
        for i, output_logits in enumerate(outputs):
            # `labels[:, i]` gets the true labels for the i-th label column
            # `field_criterions[LABEL_COLUMNS[i]]` selects the appropriate loss function
            loss += field_criterions[LABEL_COLUMNS[i]](output_logits, labels[:, i])

        optimizer.zero_grad() # Clear previous gradients
        loss.backward()       # Backpropagation
        optimizer.step()      # Update model parameters
        total_loss += loss.item() # Accumulate loss
        tqdm_loader.set_postfix(loss=loss.item()) # Update progress bar with current batch loss

    return total_loss / len(loader) # Return average loss for the epoch

def evaluate_model(model, loader):
    """
    Evaluates the given PyTorch model on a validation/test set.
    Args:
        model (torch.nn.Module): The model to evaluate.
        loader (torch.utils.data.DataLoader): DataLoader for evaluation data.
    Returns:
        tuple: A tuple containing:
            - reports (dict): Classification reports (dict format) for each label column.
            - truths (list): List of true label arrays for each label column.
            - predictions (list): List of predicted label arrays for each label column.
    """
    model.eval() # Set the model to evaluation mode (disables dropout, batch norm updates, etc.)
    # Initialize lists to store predictions and true labels for each output head
    predictions = [[] for _ in range(len(LABEL_COLUMNS))]
    truths = [[] for _ in range(len(LABEL_COLUMNS))]

    with torch.no_grad(): # Disable gradient calculations during evaluation for efficiency
        for batch in tqdm(loader, desc="Evaluation"):
            if len(batch) == 2:
                inputs, labels = batch
                input_ids = inputs['input_ids'].to(DEVICE)
                attention_mask = inputs['attention_mask'].to(DEVICE)
                labels = labels.to(DEVICE)
                outputs = model(input_ids, attention_mask)
            elif len(batch) == 3:
                inputs, metadata, labels = batch
                input_ids = inputs['input_ids'].to(DEVICE)
                attention_mask = inputs['attention_mask'].to(DEVICE)
                metadata = metadata.to(DEVICE)
                labels = labels.to(DEVICE)
                outputs = model(input_ids, attention_mask, metadata)
            else:
                raise ValueError("Unsupported batch format.")

            for i, output_logits in enumerate(outputs):
                # Get the predicted class by taking the argmax of the logits
                preds = torch.argmax(output_logits, dim=1).cpu().numpy()
                predictions[i].extend(preds)
                # Get the true labels for the current output head
                truths[i].extend(labels[:, i].cpu().numpy())

    reports = {}
    # Generate classification report for each label column
    for i, col in enumerate(LABEL_COLUMNS):
        try:
            # `zero_division=0` handles cases where a class might have no true or predicted samples
            reports[col] = classification_report(truths[i], predictions[i], output_dict=True, zero_division=0)
        except ValueError:
            # Handle cases where a label might not appear in the validation set,
            # which could cause classification_report to fail.
            print(f"Warning: Could not generate classification report for {col}. Skipping.")
            reports[col] = {'accuracy': 0, 'weighted avg': {'precision': 0, 'recall': 0, 'f1-score': 0, 'support': 0}}
    return reports, truths, predictions

def summarize_metrics(metrics):
    """
    Summarizes classification reports into a readable Pandas DataFrame.
    Args:
        metrics (dict): Dictionary of classification reports, as returned by `evaluate_model`.
    Returns:
        pd.DataFrame: A DataFrame summarizing precision, recall, f1-score, accuracy, and support for each field.
    """
    summary = []
    for field, report in metrics.items():
        # Safely get metrics, defaulting to 0 if not present (e.g., for empty reports)
        precision = report['weighted avg']['precision'] if 'weighted avg' in report else 0
        recall = report['weighted avg']['recall'] if 'weighted avg' in report else 0
        f1 = report['weighted avg']['f1-score'] if 'weighted avg' in report else 0
        support = report['weighted avg']['support'] if 'weighted avg' in report else 0
        accuracy = report['accuracy'] if 'accuracy' in report else 0 # Accuracy is usually top-level
        summary.append({
            "Field": field,
            "Precision": precision,
            "Recall": recall,
            "F1-Score": f1,
            "Accuracy": accuracy,
            "Support": support
        })
    return pd.DataFrame(summary)

def save_model(model, model_name, save_format='pth'):
    """
    Saves the state dictionary of a PyTorch model.
    Args:
        model (torch.nn.Module): The trained PyTorch model.
        model_name (str): A descriptive name for the model (used for filename).
        save_format (str): Format to save the model in ('pth' for PyTorch models, 'pickle' for traditional ML models).
    """
    # Construct the save path dynamically relative to the project root
    if save_format == 'pth':
        model_path = os.path.join(MODEL_SAVE_DIR, f"{model_name}_model.pth")
        torch.save(model.state_dict(), model_path)
    elif save_format == 'pickle':
        model_path = os.path.join(MODEL_SAVE_DIR, f"{model_name}.pkl")
        joblib.dump(model, model_path)
    else:
        raise ValueError(f"Unsupported save format: {save_format}")
    
    print(f"Model saved to {model_path}")

def load_model_state(model, model_name, model_class, num_labels, metadata_dim=0):
    """
    Loads the state dictionary into a PyTorch model.
    Args:
        model (torch.nn.Module): An initialized model instance (architecture).
        model_name (str): The name of the model to load.
        model_class (class): The class of the model (e.g., BertMultiOutputModel).
        num_labels (list): List of number of classes for each label.
        metadata_dim (int): Dimensionality of metadata features, if applicable (default 0 for text-only).
    Returns:
        torch.nn.Module: The model with loaded state_dict, moved to the correct device, and set to eval mode.
    """
    model_path = os.path.join(MODEL_SAVE_DIR, f"{model_name}_model.pth")
    if not os.path.exists(model_path):
        print(f"Warning: Model file not found at {model_path}. Returning a newly initialized model instance.")
        # Re-initialize the model if not found, to ensure it has the correct architecture
        if metadata_dim > 0:
            return model_class(num_labels, metadata_dim=metadata_dim).to(DEVICE)
        else:
            return model_class(num_labels).to(DEVICE)

    model.load_state_dict(torch.load(model_path, map_location=DEVICE))
    model.to(DEVICE)
    model.eval() # Set to evaluation mode after loading
    print(f"Model loaded from {model_path}")
    return model

def predict_probabilities(model, loader):
    """
    Generates prediction probabilities for each label for a given model.
    This is used for confidence scoring and feeding into a voting ensemble.
    Args:
        model (torch.nn.Module): The trained PyTorch model.
        loader (torch.utils.data.DataLoader): DataLoader for the data to predict on.
    Returns:
        list: A list of lists of numpy arrays. Each inner list corresponds to a label column,
              containing the softmax probabilities for each sample for that label.
    """
    model.eval() # Set to evaluation mode
    # List to store probabilities for each output head
    all_probabilities = [[] for _ in range(len(LABEL_COLUMNS))]

    with torch.no_grad():
        for batch in tqdm(loader, desc="Predicting Probabilities"):
            # Unpack batch, ignoring labels as we only need inputs
            if len(batch) == 2:
                inputs, _ = batch
                input_ids = inputs['input_ids'].to(DEVICE)
                attention_mask = inputs['attention_mask'].to(DEVICE)
                outputs = model(input_ids, attention_mask)
            elif len(batch) == 3:
                inputs, metadata, _ = batch
                input_ids = inputs['input_ids'].to(DEVICE)
                attention_mask = inputs['attention_mask'].to(DEVICE)
                metadata = metadata.to(DEVICE)
                outputs = model(input_ids, attention_mask, metadata)
            else:
                raise ValueError("Unsupported batch format.")

            for i, out_logits in enumerate(outputs):
                # Apply softmax to logits to get probabilities
                probs = torch.softmax(out_logits, dim=1).cpu().numpy()
                all_probabilities[i].extend(probs)
    return all_probabilities