Blame for rmit_hackathon_2026/train.py

"""
RMIT GenAI & Cybersecurity Hackathon 2026 -- Challenge 1
Multilingual Multi-Class Prompt-Attack Classification (3 classes)
Approved Models:
  - distilbert-base-multilingual-cased (fastest baseline)
  - jhu-clsp/mmBERT-small / mmBERT-base
  - FacebookAI/xlm-roberta-base / xlm-roberta-large (strong multilingual)
  - google/rembert
  - protectai/deberta-v3-base-prompt-injection-v2
  - patronus-studio/wolf-defender-prompt-injection
"""
import os
import gc
import argparse
import numpy as np
import pandas as pd
import torch
import emoji
from torch.utils.data import DataLoader
from sklearn.model_selection import StratifiedKFold
from sklearn.utils.class_weight import compute_class_weight
from transformers import (
    AutoTokenizer,
    AutoModelForSequenceClassification,
    AutoConfig,
    get_cosine_schedule_with_warmup,
)
from tqdm.auto import tqdm

from src.utils import seed_everything, detailed_metrics
from src.dataset import PromptDataset
from src.preprocessing import clean_and_deobfuscate_text


# ---------------------------------------------------------------------------
# Approved Models Registry (RMIT Hackathon 2026 Rules)
# ---------------------------------------------------------------------------
APPROVED_MODELS = {
    "xlm-roberta-base": "FacebookAI/xlm-roberta-base",
    "xlm-roberta-large": "FacebookAI/xlm-roberta-large",
    "distilbert-multilingual": "distilbert/distilbert-base-multilingual-cased",
    "deberta-v3-prompt-injection": "protectai/deberta-v3-base-prompt-injection-v2",
    "wolf-defender": "patronus-studio/wolf-defender-prompt-injection",
    "rembert": "google/rembert",
    "mmbert-base": "jhu-clsp/mmBERT-base",
    "mmbert-small": "jhu-clsp/mmBERT-small",
}

LABEL_MAP = {"BENIGN": 0, "JAILBREAK": 1, "PROMPT_INJECTION": 2}
ID2LABEL = {0: "BENIGN", 1: "JAILBREAK", 2: "PROMPT_INJECTION"}
LABEL2ID = {"BENIGN": 0, "JAILBREAK": 1, "PROMPT_INJECTION": 2}


def get_data_paths(data_dir: str = "./data"):
    """Tu dong tim train va test csv theo data_dir cua nguoi dung hoac Kaggle input."""
    if os.path.exists("./data/train.csv"):
        return "./data/train.csv", "./data/test.csv"
    if os.path.exists(r"D:\Code\rmit\data\Data-rmit2026\train.csv"):
        return r"D:\Code\rmit\data\Data-rmit2026\train.csv", r"D:\Code\rmit\data\Data-rmit2026\test.csv"

    train_path = os.path.join(data_dir, "train.csv")
    test_path = os.path.join(data_dir, "test.csv")
    return train_path, test_path


def run_training_2026(
    model_key: str = "xlm-roberta-base",
    data_dir: str = r"D:\Code\rmit\data\Data-rmit2026",
    epochs: int = 3,
    batch_size: int = 16,
    lr: float = 2e-5,
    max_length: int = 256,
    n_folds: int = 5,
    use_class_weights: bool = True,
    quick_run: bool = False,
):
    """
    Pipeline huan luyen da ngon ngu 3-class cho RMIT Hackathon 2026:
    - 3 classes: BENIGN (62%), JAILBREAK (31%), PROMPT_INJECTION (7%)
    - Weighted Cross-Entropy Loss de giai quyet class imbalance (nang f1 cua PROMPT_INJECTION).
    - Stratified K-Fold theo ca (label + language).
    - Tinh Macro-F1 chuan xac theo tieu chi danh gia cua Kaggle.
    """
    seed_everything(42)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model_name = APPROVED_MODELS.get(model_key, model_key)
    model_alias = model_key.replace("/", "_").replace("-", "_")

    print(f"\n{'='*70}")
    print(f"RMIT Hackathon 2026 -- Challenge 1: Multilingual Prompt-Attack")
    print(f"Model backbone : {model_name}")
    print(f"Device         : {device}")
    print(f"Folds          : {'1 (quick_run)' if quick_run else n_folds}")
    print(f"Classes        : BENIGN (0), JAILBREAK (1), PROMPT_INJECTION (2)")
    print(f"{'='*70}\n")

    train_file, test_file = get_data_paths(data_dir)
    if not os.path.exists(train_file):
        raise FileNotFoundError(f"Khong tim thay tap train tai: {train_file}")

    # 1. Load data
    df_train = pd.read_csv(train_file)
    print(f"[Data] Train samples: {len(df_train)}")
    print(f"[Data] Ngon ngu: {dict(df_train['language'].value_counts())}")

    # De-obfuscate text
    df_train["clean_text"] = df_train["text"].apply(clean_and_deobfuscate_text)
    df_train["label_num"] = df_train["label"].map(lambda l: LABEL_MAP.get(str(l).strip().upper(), 0))
    print(f"[Data] Ti le nhan:\n{df_train['label'].value_counts(normalize=True)}\n")

    # Load test if exists
    df_test = None
    if os.path.exists(test_file):
        df_test = pd.read_csv(test_file)
        df_test["clean_text"] = df_test["text"].apply(clean_and_deobfuscate_text)
        print(f"[Data] Test samples: {len(df_test)}")

    # 2. Stratified Folds (stratify theo ca label va language de fold can bang)
    df_train["strat_key"] = df_train["language"].astype(str) + "_" + df_train["label_num"].astype(str)
    skf = StratifiedKFold(n_splits=n_folds, shuffle=True, random_state=42)
    df_train["fold"] = -1
    for fold, (_, val_idx) in enumerate(skf.split(df_train, df_train["strat_key"])):
        df_train.loc[val_idx, "fold"] = fold

    # 3. Class Weights cho Loss (giai quyet bat can bang nhan 62% vs 7%)
    weights = None
    if use_class_weights:
        cls_weights = compute_class_weight(
            class_weight="balanced",
            classes=np.array([0, 1, 2]),
            y=df_train["label_num"].values,
        )
        weights = torch.tensor(cls_weights, dtype=torch.float).to(device)
        print(f"[Loss] Class weights (balanced): {cls_weights.tolist()}")

    criterion = torch.nn.CrossEntropyLoss(weight=weights)

    # 4. Tokenizer & model save directory
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    save_dir = f"./models_2026/{model_alias}"
    os.makedirs(save_dir, exist_ok=True)
    os.makedirs("./models_2026", exist_ok=True)

    oof_probs = np.zeros((len(df_train), 3))
    test_probs = np.zeros((len(df_test), 3)) if df_test is not None else None

    test_loader = None
    if df_test is not None:
        test_ds = PromptDataset(df_test["clean_text"].values, tokenizer=tokenizer, max_length=max_length)
        test_loader = DataLoader(test_ds, batch_size=batch_size * 2, shuffle=False, num_workers=0)

    num_folds = 1 if quick_run else n_folds

    # 5. Training Loop
    for fold in range(num_folds):
        print(f"\n{'─'*55}")
        print(f"FOLD {fold + 1} / {num_folds}")
        print(f"{'─'*55}")

        train_df = df_train[df_train["fold"] != fold].reset_index(drop=True)
        val_df   = df_train[df_train["fold"] == fold].reset_index(drop=True)
        val_idx  = df_train[df_train["fold"] == fold].index.values

        if quick_run:
            train_df = train_df.head(64)
            val_df   = val_df.head(32)
            val_idx  = val_idx[:32]

        train_loader = DataLoader(
            PromptDataset(train_df["clean_text"].values, train_df["label_num"].values, tokenizer, max_length),
            batch_size=batch_size, shuffle=True, num_workers=0,
        )
        val_loader = DataLoader(
            PromptDataset(val_df["clean_text"].values, val_df["label_num"].values, tokenizer, max_length),
            batch_size=batch_size * 2, shuffle=False, num_workers=0,
        )

        save_fold_path = f"{save_dir}/{model_alias}_fold{fold}.pt"
        config = AutoConfig.from_pretrained(model_name, num_labels=3, id2label=ID2LABEL, label2id=LABEL2ID)
        model = AutoModelForSequenceClassification.from_pretrained(model_name, config=config).to(device)

        # Check resume
        if os.path.exists(save_fold_path):
            print(f"[Resume] Fold {fold + 1} da ton tai checkpoint: {save_fold_path}. Dang load...")
            model.load_state_dict(torch.load(save_fold_path, map_location=device))
            model.eval()
            val_preds_list = []
            with torch.no_grad():
                for batch in val_loader:
                    logits = model(batch["input_ids"].to(device), attention_mask=batch["attention_mask"].to(device)).logits
                    probs = torch.softmax(logits, dim=-1).cpu().numpy()
                    val_preds_list.extend(probs)
            val_preds_arr = np.array(val_preds_list)
            oof_probs[val_idx] = val_preds_arr
            metrics = detailed_metrics(val_df["label_num"].values, val_preds_arr)
            print(f"[Resume] Fold {fold + 1} Loaded: Macro-F1={metrics['macro_f1']} (Benign:{metrics['f1_benign']}, Jailbreak:{metrics['f1_jailbreak']}, Injection:{metrics['f1_prompt_injection']})")

            if test_loader is not None and test_probs is not None:
                with torch.no_grad():
                    test_fold_list = []
                    for batch in test_loader:
                        logits = model(batch["input_ids"].to(device), attention_mask=batch["attention_mask"].to(device)).logits
                        test_fold_list.extend(torch.softmax(logits, dim=-1).cpu().numpy())
                    test_probs += np.array(test_fold_list) / num_folds
            del model
            torch.cuda.empty_cache()
            gc.collect()
            continue

        optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=0.01)
        total_steps = len(train_loader) * (1 if quick_run else epochs)
        scheduler = get_cosine_schedule_with_warmup(
            optimizer, num_warmup_steps=int(total_steps * 0.1), num_training_steps=total_steps
        )

        best_val_f1 = 0.0
        for epoch in range(1 if quick_run else epochs):
            model.train()
            train_loss = 0.0
            for batch in tqdm(train_loader, desc=f"  Epoch {epoch+1} Train", leave=False):
                optimizer.zero_grad()
                out = model(batch["input_ids"].to(device), attention_mask=batch["attention_mask"].to(device))
                loss = criterion(out.logits, batch["label"].to(device))
                loss.backward()
                torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
                optimizer.step()
                scheduler.step()
                train_loss += loss.item()

            # Validation
            model.eval()
            val_preds_list = []
            with torch.no_grad():
                for batch in val_loader:
                    logits = model(batch["input_ids"].to(device), attention_mask=batch["attention_mask"].to(device)).logits
                    probs = torch.softmax(logits, dim=-1).cpu().numpy()
                    val_preds_list.extend(probs)

            val_preds_arr = np.array(val_preds_list)
            metrics = detailed_metrics(val_df["label_num"].values, val_preds_arr)
            avg_loss = train_loss / len(train_loader)

            print(
                f"  Epoch {epoch+1} | Loss: {avg_loss:.4f} "
                f"| Macro-F1: {metrics['macro_f1']:.4f} "
                f"[B: {metrics['f1_benign']:.4f}, JB: {metrics['f1_jailbreak']:.4f}, PI: {metrics['f1_prompt_injection']:.4f}]"
            )

            if metrics["macro_f1"] >= best_val_f1:
                best_val_f1 = metrics["macro_f1"]
                torch.save(model.state_dict(), save_fold_path)

        oof_probs[val_idx] = val_preds_arr

        # Test inference for this fold
        if test_loader is not None and test_probs is not None:
            model.eval()
            test_fold_list = []
            with torch.no_grad():
                for batch in test_loader:
                    logits = model(batch["input_ids"].to(device), attention_mask=batch["attention_mask"].to(device)).logits
                    test_fold_list.extend(torch.softmax(logits, dim=-1).cpu().numpy())
            test_probs += np.array(test_fold_list) / num_folds

        del model, optimizer
        torch.cuda.empty_cache()
        gc.collect()

    # 6. Save OOF & Test probs
    np.save(f"./models_2026/oof_{model_alias}.npy", oof_probs)
    print(f"\n[Save] Da luu OOF probs -> ./models_2026/oof_{model_alias}.npy")

    overall_metrics = detailed_metrics(df_train["label_num"].values, oof_probs)
    print(f"[Done] Overall OOF Macro-F1 = {overall_metrics['macro_f1']:.4f}")
    print(f"       Per-Class: Benign={overall_metrics['f1_benign']:.4f}, Jailbreak={overall_metrics['f1_jailbreak']:.4f}, Injection={overall_metrics['f1_prompt_injection']:.4f}")

    if test_probs is not None:
        np.save(f"./models_2026/test_{model_alias}.npy", test_probs)
        pred_labels = [ID2LABEL[idx] for idx in np.argmax(test_probs, axis=1)]
        sub = pd.DataFrame({
            "id": df_test["id"],
            "label": pred_labels
        })
        sub_file = f"submission_{model_alias}.csv"
        sub.to_csv(sub_file, index=False)
        sub.to_csv("submission.csv", index=False)
        print(f"\n[Done] Da tao file nop bai: {sub_file} va submission.csv (6,000 dong)")
        print(f"[Done] Ti le du doan:\n{sub['label'].value_counts(normalize=True)}")

    return overall_metrics


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="RMIT Hackathon 2026 -- C1 Multilingual Classifier")
    parser.add_argument("--model", default="xlm-roberta-base", choices=list(APPROVED_MODELS.keys()),
                        help="Approved model backbone")
    parser.add_argument("--data_dir", default=r"D:\Code\rmit\data\Data-rmit2026")
    parser.add_argument("--epochs", type=int, default=3)
    parser.add_argument("--batch", type=int, default=16)
    parser.add_argument("--lr", type=float, default=2e-5)
    parser.add_argument("--folds", type=int, default=5)
    parser.add_argument("--maxlen", type=int, default=256)
    parser.add_argument("--quick", action="store_true", help="Quick run 64 samples")
    args = parser.parse_args()

    run_training_2026(
        model_key=args.model,
        data_dir=args.data_dir,
        epochs=args.epochs,
        batch_size=args.batch,
        lr=args.lr,
        max_length=args.maxlen,
        n_folds=args.folds,
        quick_run=args.quick,
    )