File size: 3,419 Bytes
4540fc1 | 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 | """General-sentiment trainer — SST-2, targeting the 3.9M-dl/mo distilbert-sst2 incumbent.
Same discipline as finsense_train.py: fp32, best-ckpt by val metric, held-out
reporting. SST-2 official: train 67,349 / validation 872 (the reported split —
incumbent's 91.3 is on validation; test is unlabeled). We carve 5% of train as
our early-stop val and report on the official validation set, untouched.
Usage: python sst2_train.py [--seed 42] [--base answerdotai/ModernBERT-base]
"""
import argparse
import json
import os
import torch
from datasets import load_dataset
from sklearn.metrics import accuracy_score, f1_score
from transformers import (AutoModelForSequenceClassification, AutoTokenizer,
Trainer, TrainingArguments)
p = argparse.ArgumentParser()
p.add_argument("--seed", type=int, default=42)
p.add_argument("--base", default="answerdotai/ModernBERT-base")
p.add_argument("--epochs", type=int, default=2)
args = p.parse_args()
OUT = os.path.expanduser(f"~/finsense_runs/sst2-{args.base.split('/')[-1]}-s{args.seed}")
os.makedirs(OUT, exist_ok=True)
LABELS = ["negative", "positive"]
ds = load_dataset("nyu-mll/glue", "sst2")
full_train = ds["train"].shuffle(seed=42)
n_val = int(len(full_train) * 0.05)
early_val = full_train.select(range(n_val))
train = full_train.select(range(n_val, len(full_train)))
report_val = ds["validation"] # incumbent's benchmark split — never used for selection
print(f"train={len(train)} early_val={len(early_val)} report_val={len(report_val)}", flush=True)
tok = AutoTokenizer.from_pretrained(args.base)
def enc(d):
return d.map(lambda b: tok(b["sentence"], truncation=True, max_length=128),
batched=True, remove_columns=[c for c in d.column_names if c not in ("label",)])
torch.manual_seed(args.seed)
model = AutoModelForSequenceClassification.from_pretrained(
args.base, num_labels=2,
id2label=dict(enumerate(LABELS)), label2id={l: i for i, l in enumerate(LABELS)})
def metrics(pred):
y, yhat = pred.label_ids, pred.predictions.argmax(-1)
return {"accuracy": accuracy_score(y, yhat)}
import transformers as _tf
_tok_kw = {"processing_class": tok} if int(_tf.__version__.split(".")[0]) >= 5 else {"tokenizer": tok}
trainer = Trainer(
model=model, **_tok_kw,
train_dataset=enc(train), eval_dataset=enc(early_val),
compute_metrics=metrics,
args=TrainingArguments(
output_dir=f"{OUT}/ckpt",
num_train_epochs=args.epochs,
per_device_train_batch_size=32, per_device_eval_batch_size=64,
learning_rate=2e-5, warmup_ratio=0.06, weight_decay=0.01,
eval_strategy="steps", eval_steps=500,
save_strategy="steps", save_steps=500, save_total_limit=1,
load_best_model_at_end=True, metric_for_best_model="accuracy",
logging_steps=100, seed=args.seed, report_to=[],
fp16=False, bf16=False,
),
)
trainer.train()
pred = trainer.predict(enc(report_val))
y, yhat = pred.label_ids, pred.predictions.argmax(-1)
res = {"base": args.base, "seed": args.seed, "task": "sst2",
"val_accuracy": float(accuracy_score(y, yhat)),
"val_f1": float(f1_score(y, yhat)),
"incumbent": "distilbert-sst2 = 0.913 on this same validation split"}
json.dump(res, open(f"{OUT}/result.json", "w"), indent=1)
print("SST2_RESULT", json.dumps(res), flush=True)
model.save_pretrained(f"{OUT}/model")
tok.save_pretrained(f"{OUT}/model")
|