| """ |
| 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 = { |
| "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}") |
|
|
| |
| 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())}") |
|
|
| |
| 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") |
|
|
| |
| 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)}") |
|
|
| |
| 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 |
|
|
| |
| 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) |
|
|
| |
| 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 |
|
|
| |
| 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) |
|
|
| |
| 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() |
|
|
| |
| 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 |
|
|
| |
| 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() |
|
|
| |
| 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, |
| ) |
|
|
