Snaseem2026/code-comment-classifier

1"""2Utility functions for training and evaluation3"""4import numpy as np5from sklearn.metrics import (6    accuracy_score, 7    precision_recall_fscore_support, 8    confusion_matrix,9    classification_report10)11import matplotlib.pyplot as plt12import seaborn as sns13from typing import Dict, Tuple, List, Optional14import os15 16 17def compute_metrics(eval_pred, id2label: Optional[Dict[int, str]] = None) -> Dict[str, float]:18    """19    Compute comprehensive metrics for evaluation.20    21    Args:22        eval_pred: Tuple of (predictions, labels)23        id2label: Optional mapping from label IDs to label names for per-class metrics24        25    Returns:26        Dictionary of metrics including overall and per-class metrics27    """28    predictions, labels = eval_pred29    predictions = np.argmax(predictions, axis=1)30    31    # Overall metrics32    accuracy = accuracy_score(labels, predictions)33    34    # Weighted metrics (accounts for class imbalance)35    precision_weighted, recall_weighted, f1_weighted, _ = precision_recall_fscore_support(36        labels,37        predictions,38        average='weighted',39        zero_division=040    )41    42    # Macro-averaged metrics (treats all classes equally)43    precision_macro, recall_macro, f1_macro, _ = precision_recall_fscore_support(44        labels,45        predictions,46        average='macro',47        zero_division=048    )49    50    # Micro-averaged metrics (aggregates contributions of all classes)51    precision_micro, recall_micro, f1_micro, _ = precision_recall_fscore_support(52        labels,53        predictions,54        average='micro',55        zero_division=056    )57    58    metrics = {59        'accuracy': accuracy,60        'precision_weighted': precision_weighted,61        'recall_weighted': recall_weighted,62        'f1_weighted': f1_weighted,63        'precision_macro': precision_macro,64        'recall_macro': recall_macro,65        'f1_macro': f1_macro,66        'precision_micro': precision_micro,67        'recall_micro': recall_micro,68        'f1_micro': f1_micro,69    }70    71    # Per-class metrics if label mapping is provided72    if id2label is not None:73        num_classes = len(id2label)74        precision_per_class, recall_per_class, f1_per_class, support = precision_recall_fscore_support(75            labels,76            predictions,77            labels=list(range(num_classes)),78            average=None,79            zero_division=080        )81        82        for i in range(num_classes):83            label_name = id2label[i]84            metrics[f'precision_{label_name}'] = float(precision_per_class[i])85            metrics[f'recall_{label_name}'] = float(recall_per_class[i])86            metrics[f'f1_{label_name}'] = float(f1_per_class[i])87            metrics[f'support_{label_name}'] = int(support[i])88    89    return metrics90 91 92def compute_metrics_factory(id2label: Optional[Dict[int, str]] = None):93    """94    Factory function to create compute_metrics with label mapping.95    96    Args:97        id2label: Mapping from label IDs to label names98        99    Returns:100        Function compatible with HuggingFace Trainer101    """102    def compute_metrics_fn(eval_pred):103        return compute_metrics(eval_pred, id2label)104    105    return compute_metrics_fn106 107 108def plot_confusion_matrix(109    y_true: np.ndarray,110    y_pred: np.ndarray,111    labels: List[str],112    save_path: str = "confusion_matrix.png",113    normalize: bool = False,114    figsize: Tuple[int, int] = (10, 8)115) -> None:116    """117    Plot and save confusion matrix with optional normalization.118    119    Args:120        y_true: True labels121        y_pred: Predicted labels122        labels: List of label names123        save_path: Path to save the plot124        normalize: If True, normalize confusion matrix to percentages125        figsize: Figure size (width, height)126    """127    cm = confusion_matrix(y_true, y_pred, labels=list(range(len(labels))))128    129    if normalize:130        cm = cm.astype('float') / cm.sum(axis=1)[:, np.newaxis]131        fmt = '.2f'132        title = 'Normalized Confusion Matrix'133    else:134        fmt = 'd'135        title = 'Confusion Matrix'136    137    plt.figure(figsize=figsize)138    sns.heatmap(139        cm,140        annot=True,141        fmt=fmt,142        cmap='Blues',143        xticklabels=labels,144        yticklabels=labels,145        cbar_kws={'label': 'Percentage' if normalize else 'Count'}146    )147    plt.title(title, fontsize=14, fontweight='bold')148    plt.ylabel('True Label', fontsize=12)149    plt.xlabel('Predicted Label', fontsize=12)150    plt.tight_layout()151    152    # Create directory if it doesn't exist153    os.makedirs(os.path.dirname(save_path) if os.path.dirname(save_path) else '.', exist_ok=True)154    plt.savefig(save_path, dpi=300, bbox_inches='tight')155    plt.close()156    157    print(f"Confusion matrix saved to {save_path}")158 159 160def print_classification_report(161    y_true: np.ndarray,162    y_pred: np.ndarray,163    labels: List[str],164    output_dict: bool = False165) -> Optional[Dict]:166    """167    Print detailed classification report.168    169    Args:170        y_true: True labels171        y_pred: Predicted labels172        labels: List of label names173        output_dict: If True, return report as dictionary instead of printing174        175    Returns:176        Classification report as dictionary if output_dict=True, else None177    """178    report = classification_report(179        y_true,180        y_pred,181        target_names=labels,182        digits=4,183        output_dict=output_dict,184        zero_division=0185    )186    187    if output_dict:188        return report189    190    print("\nClassification Report:")191    print("=" * 60)192    print(report)193    return None194 195 196def plot_training_curves(197    train_losses: List[float],198    eval_losses: List[float],199    eval_metrics: Dict[str, List[float]],200    save_path: str = "./results/training_curves.png"201) -> None:202    """203    Plot training and evaluation curves.204    205    Args:206        train_losses: List of training losses per step/epoch207        eval_losses: List of evaluation losses per step/epoch208        eval_metrics: Dictionary of metric names to lists of values209        save_path: Path to save the plot210    """211    fig, axes = plt.subplots(2, 2, figsize=(15, 10))212    213    # Loss curves214    axes[0, 0].plot(train_losses, label='Train Loss', color='blue')215    axes[0, 0].plot(eval_losses, label='Eval Loss', color='red')216    axes[0, 0].set_xlabel('Step/Epoch')217    axes[0, 0].set_ylabel('Loss')218    axes[0, 0].set_title('Training and Validation Loss')219    axes[0, 0].legend()220    axes[0, 0].grid(True, alpha=0.3)221    222    # Accuracy223    if 'accuracy' in eval_metrics:224        axes[0, 1].plot(eval_metrics['accuracy'], label='Accuracy', color='green')225        axes[0, 1].set_xlabel('Step/Epoch')226        axes[0, 1].set_ylabel('Accuracy')227        axes[0, 1].set_title('Validation Accuracy')228        axes[0, 1].legend()229        axes[0, 1].grid(True, alpha=0.3)230    231    # F1 Score232    if 'f1_weighted' in eval_metrics:233        axes[1, 0].plot(eval_metrics['f1_weighted'], label='F1 (weighted)', color='purple')234        axes[1, 0].set_xlabel('Step/Epoch')235        axes[1, 0].set_ylabel('F1 Score')236        axes[1, 0].set_title('Validation F1 Score')237        axes[1, 0].legend()238        axes[1, 0].grid(True, alpha=0.3)239    240    # Precision and Recall241    if 'precision_weighted' in eval_metrics and 'recall_weighted' in eval_metrics:242        axes[1, 1].plot(eval_metrics['precision_weighted'], label='Precision', color='orange')243        axes[1, 1].plot(eval_metrics['recall_weighted'], label='Recall', color='cyan')244        axes[1, 1].set_xlabel('Step/Epoch')245        axes[1, 1].set_ylabel('Score')246        axes[1, 1].set_title('Validation Precision and Recall')247        axes[1, 1].legend()248        axes[1, 1].grid(True, alpha=0.3)249    250    plt.tight_layout()251    os.makedirs(os.path.dirname(save_path) if os.path.dirname(save_path) else '.', exist_ok=True)252    plt.savefig(save_path, dpi=300, bbox_inches='tight')253    plt.close()254    255    print(f"Training curves saved to {save_path}")256