| """ |
| Utility module -- Member A (Huy) |
| Seed, metrics, and helper functions for model training. |
| """ |
| import os |
| import random |
| import numpy as np |
| import torch |
| from sklearn.metrics import f1_score, roc_auc_score, precision_score, recall_score, confusion_matrix |
|
|
|
|
| def seed_everything(seed: int = 42) -> None: |
| """Fix tat ca random seeds de ket qua co the tai tao.""" |
| random.seed(seed) |
| os.environ["PYTHONHASHSEED"] = str(seed) |
| np.random.seed(seed) |
| torch.manual_seed(seed) |
| torch.cuda.manual_seed(seed) |
| torch.cuda.manual_seed_all(seed) |
| torch.backends.cudnn.deterministic = True |
| torch.backends.cudnn.benchmark = False |
|
|
|
|
| def compute_metrics(eval_pred) -> dict: |
| """ |
| Tinh Macro-F1 tu logits va labels cho 3 classes (BENIGN, JAILBREAK, PROMPT_INJECTION). |
| eval_pred: tuple (logits, labels) tuong thich voi HuggingFace Trainer. |
| """ |
| predictions, labels = eval_pred |
| if isinstance(predictions, tuple): |
| predictions = predictions[0] |
| preds = np.argmax(predictions, axis=1) |
|
|
| macro_f1 = f1_score(labels, preds, average="macro", zero_division=0) |
| per_class_f1 = f1_score(labels, preds, average=None, zero_division=0) |
| return { |
| "macro_f1": macro_f1, |
| "f1_benign": per_class_f1[0] if len(per_class_f1) > 0 else 0.0, |
| "f1_jailbreak": per_class_f1[1] if len(per_class_f1) > 1 else 0.0, |
| "f1_prompt_injection": per_class_f1[2] if len(per_class_f1) > 2 else 0.0, |
| } |
|
|
|
|
| def detailed_metrics(y_true: np.ndarray, y_pred_probs: np.ndarray) -> dict: |
| """ |
| Bao cao day du cac chi so phan loai 3-class (Macro-F1, Precision, Recall per class). |
| y_pred_probs: mang xac suat shape (N, 3) hoac logits (N, 3). |
| """ |
| if y_pred_probs.ndim == 2: |
| y_pred = np.argmax(y_pred_probs, axis=1) |
| else: |
| y_pred = y_pred_probs |
|
|
| macro_f1 = f1_score(y_true, y_pred, average="macro", zero_division=0) |
| per_class_f1 = f1_score(y_true, y_pred, average=None, zero_division=0) |
| precision = precision_score(y_true, y_pred, average="macro", zero_division=0) |
| recall = recall_score(y_true, y_pred, average="macro", zero_division=0) |
| cm = confusion_matrix(y_true, y_pred) |
| return { |
| "macro_f1": round(float(macro_f1), 4), |
| "f1_benign": round(float(per_class_f1[0]), 4) if len(per_class_f1) > 0 else 0.0, |
| "f1_jailbreak": round(float(per_class_f1[1]), 4) if len(per_class_f1) > 1 else 0.0, |
| "f1_prompt_injection": round(float(per_class_f1[2]), 4) if len(per_class_f1) > 2 else 0.0, |
| "precision": round(float(precision), 4), |
| "recall": round(float(recall), 4), |
| "confusion_matrix": cm.tolist(), |
| } |
|
|
