File size: 3,399 Bytes
bc3ba9e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 | # EmpathRAG — DeBERTa NLI Safety Guardrail Fine-Tuning
# Run on Google Colab Pro (A100). Expected time: under 2 hours.
# Target: recall > 0.80, precision > 0.65
#
# SETUP INSTRUCTIONS:
# 1. Upload nli_train.csv, nli_val.csv, nli_test.csv to Colab (or mount Drive)
# 2. Set runtime to A100 GPU
# 3. Run all cells in order
# ── Cell 1: Install ──────────────────────────────────────────────────────────
# !pip install transformers datasets evaluate scikit-learn accelerate -q
# ── Cell 2: Mount Drive ──────────────────────────────────────────────────────
# from google.colab import drive
# drive.mount("/content/drive")
# SAVE_DIR = "/content/drive/MyDrive/empathrag/safety_guardrail"
# !mkdir -p {SAVE_DIR}
# ── Cell 3: Training script ──────────────────────────────────────────────────
import pandas as pd
import numpy as np
import torch
import evaluate as evaluate_lib
from datasets import Dataset
from transformers import (
AutoTokenizer,
AutoModelForSequenceClassification,
TrainingArguments,
Trainer,
)
SAVE_DIR = "/content/drive/MyDrive/empathrag/safety_guardrail"
train_df = pd.read_csv("nli_train.csv")
val_df = pd.read_csv("nli_val.csv")
test_df = pd.read_csv("nli_test.csv")
MODEL = "microsoft/deberta-v3-base"
tokenizer = AutoTokenizer.from_pretrained(MODEL)
def tokenize(batch):
return tokenizer(
batch["text"],
batch["hypothesis"],
truncation=True,
max_length=256,
padding="max_length",
)
train_ds = Dataset.from_pandas(train_df).map(tokenize, batched=True)
val_ds = Dataset.from_pandas(val_df).map(tokenize, batched=True)
test_ds = Dataset.from_pandas(test_df).map(tokenize, batched=True)
model = AutoModelForSequenceClassification.from_pretrained(MODEL, num_labels=2)
f1 = evaluate_lib.load("f1")
recall = evaluate_lib.load("recall")
precision = evaluate_lib.load("precision")
def compute_metrics(eval_pred):
logits, labels = eval_pred
preds = np.argmax(logits, axis=-1)
return {
"f1": f1.compute(predictions=preds, references=labels, pos_label=0)["f1"],
"recall": recall.compute(predictions=preds, references=labels, pos_label=0)["recall"],
"precision": precision.compute(predictions=preds, references=labels, pos_label=0)["precision"],
}
args = TrainingArguments(
output_dir=SAVE_DIR,
num_train_epochs=4,
per_device_train_batch_size=32,
per_device_eval_batch_size=64,
learning_rate=1e-5,
warmup_ratio=0.1,
weight_decay=0.01,
eval_strategy="epoch",
save_strategy="epoch",
load_best_model_at_end=True,
metric_for_best_model="recall",
fp16=True,
report_to="none",
)
trainer = Trainer(
model=model,
args=args,
train_dataset=train_ds,
eval_dataset=val_ds,
compute_metrics=compute_metrics,
)
trainer.train()
trainer.save_model(SAVE_DIR)
tokenizer.save_pretrained(SAVE_DIR)
results = trainer.evaluate(test_ds)
print(f"Test recall: {results['eval_recall']:.3f} | precision: {results['eval_precision']:.3f}")
print("Target: recall > 0.80, precision > 0.65")
|