File size: 3,691 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 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 | # EmpathRAG — RoBERTa Emotion Classifier Fine-Tuning
# Run on Google Colab Pro (A100). Expected time: under 90 minutes.
# Target: weighted F1 > 0.75 on the 5-class taxonomy.
#
# SETUP INSTRUCTIONS:
# 1. Upload this file to Google Colab
# 2. Set runtime to A100 GPU
# 3. Run all cells in order
# ── Cell 1: Install dependencies ────────────────────────────────────────────
# !pip install transformers==4.38.2 datasets peft evaluate scikit-learn accelerate -q
# ── Cell 2: Mount Drive ──────────────────────────────────────────────────────
# from google.colab import drive
# drive.mount("/content/drive")
# SAVE_DIR = "/content/drive/MyDrive/empathrag/emotion_classifier"
# !mkdir -p {SAVE_DIR}
# ── Cell 3: Training script ──────────────────────────────────────────────────
from datasets import load_dataset
from transformers import (
AutoTokenizer,
AutoModelForSequenceClassification,
TrainingArguments,
Trainer,
)
from peft import get_peft_model, LoraConfig, TaskType
import evaluate
import numpy as np
import torch
SAVE_DIR = "/content/drive/MyDrive/empathrag/emotion_classifier"
LABEL_MAP = {
"grief": 0, "remorse": 0, "fear": 0, "sadness": 0,
"nervousness": 1, "confusion": 1, "embarrassment": 1,
"anger": 2, "annoyance": 2, "disappointment": 2, "disgust": 2,
"neutral": 3,
"optimism": 4, "relief": 4, "gratitude": 4, "joy": 4,
"love": 4, "admiration": 4, "amusement": 4, "approval": 4,
"caring": 4, "curiosity": 4, "desire": 4, "excitement": 4,
"pride": 4, "realization": 4, "surprise": 4,
}
raw = load_dataset("google-research-datasets/go_emotions", "simplified")
feature_names = raw["train"].features["labels"].feature.names
def remap(example):
coarse = 3
for lid in example["labels"]:
name = feature_names[lid]
if name in LABEL_MAP:
coarse = LABEL_MAP[name]
break
return {"label": coarse}
dataset = raw.map(remap)
MODEL = "roberta-base"
tokenizer = AutoTokenizer.from_pretrained(MODEL)
def tokenize(batch):
return tokenizer(batch["text"], truncation=True, max_length=128, padding="max_length")
tokenized = dataset.map(tokenize, batched=True)
lora_cfg = LoraConfig(
task_type=TaskType.SEQ_CLS,
r=16,
lora_alpha=32,
lora_dropout=0.1,
target_modules=["query", "value"],
)
base = AutoModelForSequenceClassification.from_pretrained(MODEL, num_labels=5)
model = get_peft_model(base, lora_cfg)
model.print_trainable_parameters()
f1_metric = evaluate.load("f1")
def compute_metrics(eval_pred):
logits, labels = eval_pred
preds = np.argmax(logits, axis=-1)
return f1_metric.compute(predictions=preds, references=labels, average="weighted")
args = TrainingArguments(
output_dir=SAVE_DIR,
num_train_epochs=5,
per_device_train_batch_size=64,
per_device_eval_batch_size=128,
learning_rate=2e-4,
warmup_ratio=0.1,
weight_decay=0.01,
eval_strategy="epoch",
save_strategy="epoch",
load_best_model_at_end=True,
metric_for_best_model="f1",
fp16=True,
report_to="none",
)
trainer = Trainer(
model=model,
args=args,
train_dataset=tokenized["train"],
eval_dataset=tokenized["validation"],
compute_metrics=compute_metrics,
)
trainer.train()
trainer.save_model(SAVE_DIR)
tokenizer.save_pretrained(SAVE_DIR)
print("Training complete — checkpoint saved to Drive")
|