Snaseem2026/code-comment-classifier

1"""2Evaluation script for trained model with comprehensive analysis3"""4import argparse5import sys6import os7import numpy as np8import pandas as pd9from transformers import AutoTokenizer, AutoModelForSequenceClassification, Trainer10 11# Add parent directory to path12sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))13 14from src import (15    load_config, 16    compute_metrics_factory, 17    plot_confusion_matrix, 18    print_classification_report19)20from src.data_loader import prepare_datasets_for_training21 22 23def analyze_errors(24    test_dataset,25    predictions: np.ndarray,26    labels: np.ndarray,27    id2label: dict,28    tokenizer,29    top_n: int = 1030) -> pd.DataFrame:31    """32    Analyze misclassified examples.33    34    Args:35        test_dataset: Test dataset36        predictions: Predicted labels37        labels: True labels38        id2label: Label mapping39        tokenizer: Tokenizer to decode text40        top_n: Number of examples to show per error type41        42    Returns:43        DataFrame with error analysis44    """45    errors = []46    for i, (pred, true_label) in enumerate(zip(predictions, labels)):47        if pred != true_label:48            # Decode the comment (approximate, as original text is removed)49            # Note: This is a limitation - we'd need to keep original text50            errors.append({51                'index': i,52                'true_label': id2label[true_label],53                'predicted_label': id2label[pred],54                'error_type': f"{id2label[true_label]} -> {id2label[pred]}"55            })56    57    error_df = pd.DataFrame(errors)58    if len(error_df) > 0:59        print(f"\nError Analysis:")60        print(f"Total errors: {len(error_df)}")61        print(f"\nError type distribution:")62        print(error_df['error_type'].value_counts())63    64    return error_df65 66 67def evaluate_model(68    model_path: str, 69    config_path: str = "config.yaml",70    save_plots: bool = True71):72    """73    Evaluate trained model on test set with comprehensive analysis.74    75    Args:76        model_path: Path to the trained model77        config_path: Path to configuration file78        save_plots: Whether to save visualization plots79    """80    print("=" * 60)81    print("Model Evaluation")82    print("=" * 60)83    84    # Load config85    config = load_config(config_path)86    87    # Create output directory88    output_dir = config['training'].get('output_dir', './results')89    os.makedirs(output_dir, exist_ok=True)90    91    # Load datasets92    print("\n[1/5] Loading datasets...")93    tokenized_datasets, label2id, id2label, _ = prepare_datasets_for_training(config_path)94    test_dataset = tokenized_datasets['test']95    print(f"✓ Test samples: {len(test_dataset)}")96    97    # Load model and tokenizer98    print("\n[2/5] Loading trained model...")99    tokenizer = AutoTokenizer.from_pretrained(model_path)100    model = AutoModelForSequenceClassification.from_pretrained(model_path)101    print(f"✓ Model loaded from {model_path}")102    103    # Create trainer for evaluation104    print("\n[3/5] Running evaluation...")105    compute_metrics_fn = compute_metrics_factory(id2label)106    trainer = Trainer(107        model=model,108        tokenizer=tokenizer,109        compute_metrics=compute_metrics_fn110    )111    112    # Get predictions113    predictions_output = trainer.predict(test_dataset)114    predictions = np.argmax(predictions_output.predictions, axis=1)115    labels = predictions_output.label_ids116    117    # Print metrics118    print("\n[4/5] Computing detailed metrics...")119    print("\n" + "=" * 60)120    print("Test Set Results")121    print("=" * 60)122    123    metrics = predictions_output.metrics124    125    # Overall metrics126    print("\nOverall Metrics:")127    overall_metrics = ['accuracy', 'f1_weighted', 'f1_macro', 'precision_weighted', 'recall_weighted']128    for metric in overall_metrics:129        key = f'test_{metric}'130        if key in metrics:131            print(f"  {metric.replace('_', ' ').title()}: {metrics[key]:.4f}")132    133    # Per-class metrics134    print("\nPer-Class Metrics:")135    label_names = [id2label[i] for i in range(len(id2label))]136    for label_name in label_names:137        precision_key = f'test_precision_{label_name}'138        recall_key = f'test_recall_{label_name}'139        f1_key = f'test_f1_{label_name}'140        if precision_key in metrics:141            print(f"\n  {label_name.upper()}:")142            print(f"    Precision: {metrics[precision_key]:.4f}")143            print(f"    Recall:    {metrics[recall_key]:.4f}")144            print(f"    F1-Score:  {metrics[f1_key]:.4f}")145            print(f"    Support:   {metrics.get(f'test_support_{label_name}', 'N/A')}")146    147    # Detailed classification report148    print("\n" + "=" * 60)149    print_classification_report(labels, predictions, label_names)150    151    # Plot confusion matrix152    print("\n[5/5] Generating visualizations...")153    if save_plots:154        plot_confusion_matrix(155            labels,156            predictions,157            label_names,158            save_path=os.path.join(output_dir, "confusion_matrix.png"),159            normalize=False160        )161        162        # Also save normalized version163        plot_confusion_matrix(164            labels,165            predictions,166            label_names,167            save_path=os.path.join(output_dir, "confusion_matrix_normalized.png"),168            normalize=True169        )170    171    # Error analysis172    error_df = analyze_errors(test_dataset, predictions, labels, id2label, tokenizer)173    if len(error_df) > 0 and save_plots:174        error_path = os.path.join(output_dir, "error_analysis.csv")175        error_df.to_csv(error_path, index=False)176        print(f"✓ Error analysis saved to {error_path}")177    178    print("\n" + "=" * 60)179    print("Evaluation Complete! 🎉")180    print("=" * 60)181    print(f"\nResults saved to: {output_dir}")182 183 184if __name__ == "__main__":185    parser = argparse.ArgumentParser(description="Evaluate trained model")186    parser.add_argument(187        "--model-path",188        type=str,189        default="./results/final_model",190        help="Path to the trained model"191    )192    parser.add_argument(193        "--config",194        type=str,195        default="config.yaml",196        help="Path to configuration file"197    )198    parser.add_argument(199        "--no-plots",200        action="store_true",201        help="Skip generating visualization plots"202    )203    args = parser.parse_args()204    205    evaluate_model(args.model_path, args.config, save_plots=not args.no_plots)206