Spaces:
Sleeping
Sleeping
| """ | |
| Model evaluation module: | |
| Computes comprehensive multi-class metrics, security-oriented performance indices, | |
| confusion matrices, and visualization plots. | |
| """ | |
| import os | |
| import matplotlib | |
| matplotlib.use("Agg") # Headless backend | |
| import matplotlib.pyplot as plt | |
| import seaborn as sns | |
| import numpy as np | |
| import pandas as pd | |
| from typing import Dict, Any, List | |
| from sklearn.metrics import ( | |
| accuracy_score, precision_score, recall_score, f1_score, | |
| classification_report, confusion_matrix, roc_auc_score | |
| ) | |
| from src.data_loader import CLASS_NAMES, ID_TO_CLASS | |
| def evaluate_predictions( | |
| y_true: np.ndarray, | |
| y_pred: np.ndarray, | |
| y_proba: np.ndarray = None, | |
| dataset_name: str = "validation" | |
| ) -> Dict[str, Any]: | |
| """ | |
| Computes rigorous evaluation metrics with a focus on security metrics. | |
| For email threats, false negatives on 'Suspicious' can lead to breaches. | |
| """ | |
| acc = accuracy_score(y_true, y_pred) | |
| macro_p = precision_score(y_true, y_pred, average="macro", zero_division=0) | |
| macro_r = recall_score(y_true, y_pred, average="macro", zero_division=0) | |
| macro_f1 = f1_score(y_true, y_pred, average="macro", zero_division=0) | |
| weighted_p = precision_score(y_true, y_pred, average="weighted", zero_division=0) | |
| weighted_r = recall_score(y_true, y_pred, average="weighted", zero_division=0) | |
| weighted_f1 = f1_score(y_true, y_pred, average="weighted", zero_division=0) | |
| # Per-class metrics | |
| class_p = precision_score(y_true, y_pred, average=None, zero_division=0) | |
| class_r = recall_score(y_true, y_pred, average=None, zero_division=0) | |
| class_f1 = f1_score(y_true, y_pred, average=None, zero_division=0) | |
| cm = confusion_matrix(y_true, y_pred) | |
| report_str = classification_report( | |
| y_true, y_pred, target_names=CLASS_NAMES, digits=4, zero_division=0 | |
| ) | |
| metrics = { | |
| f"{dataset_name}_accuracy": float(acc), | |
| f"{dataset_name}_macro_precision": float(macro_p), | |
| f"{dataset_name}_macro_recall": float(macro_r), | |
| f"{dataset_name}_macro_f1": float(macro_f1), | |
| f"{dataset_name}_weighted_f1": float(weighted_f1), | |
| } | |
| # Add per-class specific metrics | |
| for idx, name in enumerate(CLASS_NAMES): | |
| clean_name = name.lower() | |
| metrics[f"{dataset_name}_{clean_name}_precision"] = float(class_p[idx]) | |
| metrics[f"{dataset_name}_{clean_name}_recall"] = float(class_r[idx]) | |
| metrics[f"{dataset_name}_{clean_name}_f1"] = float(class_f1[idx]) | |
| # Security false negative rate for Suspicious class (class index 2) | |
| suspicious_idx = 2 | |
| suspicious_fnr = 1.0 - float(class_r[suspicious_idx]) | |
| metrics[f"{dataset_name}_suspicious_false_negative_rate"] = float(suspicious_fnr) | |
| # Multi-class ROC-AUC if probabilities provided | |
| if y_proba is not None: | |
| try: | |
| auc = roc_auc_score(y_true, y_proba, multi_class="ovr", average="macro") | |
| metrics[f"{dataset_name}_roc_auc_ovr"] = float(auc) | |
| except Exception: | |
| metrics[f"{dataset_name}_roc_auc_ovr"] = 0.0 | |
| return { | |
| "metrics": metrics, | |
| "confusion_matrix": cm, | |
| "classification_report": report_str | |
| } | |
| def plot_confusion_matrix( | |
| cm: np.ndarray, | |
| model_name: str, | |
| output_path: str | |
| ) -> str: | |
| """ | |
| Renders and saves a high-resolution confusion matrix heatmap. | |
| """ | |
| os.makedirs(os.path.dirname(output_path), exist_ok=True) | |
| plt.figure(figsize=(7, 6)) | |
| # Normalize for percentage annotations | |
| cm_norm = cm.astype('float') / cm.sum(axis=1)[:, np.newaxis] | |
| annotations = np.empty_like(cm, dtype=object) | |
| for i in range(cm.shape[0]): | |
| for j in range(cm.shape[1]): | |
| annotations[i, j] = f"{cm[i, j]}\n({cm_norm[i, j]*100:.1f}%)" | |
| sns.heatmap( | |
| cm, | |
| annot=annotations, | |
| fmt="", | |
| cmap="Blues", | |
| xticklabels=CLASS_NAMES, | |
| yticklabels=CLASS_NAMES, | |
| cbar=True, | |
| annot_kws={"size": 11, "weight": "bold"} | |
| ) | |
| plt.title(f"Confusion Matrix: {model_name}", fontsize=13, pad=12, weight="bold") | |
| plt.xlabel("Predicted Class", fontsize=11, weight="bold") | |
| plt.ylabel("True Class", fontsize=11, weight="bold") | |
| plt.tight_layout() | |
| plt.savefig(output_path, dpi=200) | |
| plt.close() | |
| return output_path | |
| def plot_metric_comparisons( | |
| comparison_df: pd.DataFrame, | |
| output_path: str | |
| ) -> str: | |
| """ | |
| Plots a multi-bar chart comparing model performances across key metrics. | |
| """ | |
| os.makedirs(os.path.dirname(output_path), exist_ok=True) | |
| plt.figure(figsize=(10, 6)) | |
| # Melt dataframe for seaborn barplot | |
| metrics_to_plot = ["Macro F1", "Suspicious Recall", "Accuracy", "Weighted F1"] | |
| melted = comparison_df.melt( | |
| id_vars=["Model"], | |
| value_vars=[m for m in metrics_to_plot if m in comparison_df.columns], | |
| var_name="Metric", | |
| value_name="Score" | |
| ) | |
| sns.barplot( | |
| data=melted, | |
| x="Metric", | |
| y="Score", | |
| hue="Model", | |
| palette="crest" | |
| ) | |
| plt.title("Model Performance Benchmark Comparison", fontsize=13, weight="bold", pad=12) | |
| plt.ylim(0.0, 1.05) | |
| plt.ylabel("Score (0.0 - 1.0)", fontsize=11, weight="bold") | |
| plt.xlabel("Evaluation Metric", fontsize=11, weight="bold") | |
| plt.legend(loc="lower right") | |
| plt.grid(axis="y", linestyle="--", alpha=0.5) | |
| plt.tight_layout() | |
| plt.savefig(output_path, dpi=200) | |
| plt.close() | |
| return output_path | |