File size: 6,993 Bytes
b18fe8f | 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 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 | """Train one of the five models on the graphics card and write its test scores (2026-10-29).
python models/train.py <name> e.g. business-pii-detector
Reads models/<name>/data/{train,test}.jsonl (models/build.py), fine-tunes distilroberta-base, saves the model to
models/<name>/out/ (not committed; uploaded to Hugging Face) and the scores to models/<name>/scores.json.
Span models (rows with "entities") are scored by exact entity match; classifiers by accuracy and macro F1.
"""
import json
import os
import sys
from collections import Counter
from pathlib import Path
import torch
from torch.utils.data import DataLoader
from transformers import (AutoModelForSequenceClassification, AutoModelForTokenClassification, AutoTokenizer,
get_linear_schedule_with_warmup)
BASE = "distilroberta-base"
ROOT = Path(__file__).resolve().parent
name = sys.argv[1]
folder = ROOT / name
read = lambda s: [json.loads(l) for l in open(folder / "data" / f"{s}.jsonl", encoding="utf-8")]
train_rows, test_rows = read("train"), read("test")
spans = "entities" in train_rows[0]
dev = "cuda" if torch.cuda.is_available() else "cpu"
torch.manual_seed(1029)
tok = AutoTokenizer.from_pretrained(BASE, add_prefix_space=True)
if spans:
types = sorted({e["label"] for r in train_rows + test_rows for e in r["entities"]})
labels = ["O"] + [f"{p}-{t}" for t in types for p in ("B", "I")]
else:
labels = sorted({r["label"] for r in train_rows + test_rows})
lid = {l: i for i, l in enumerate(labels)}
MAX = 256 if not spans else 160
def encode_spans(r):
enc = tok(r["text"], truncation=True, max_length=MAX, return_offsets_mapping=True)
tags = []
for s, e in enc["offset_mapping"]:
if s == e:
tags.append(-100)
continue
tag = "O"
for ent in r["entities"]:
if s >= ent["start"] and e <= ent["end"] + 1 and s < ent["end"]:
tag = ("B-" if s <= ent["start"] else "I-") + ent["label"]
break
tags.append(lid[tag])
enc.pop("offset_mapping")
enc["labels"] = tags
return enc
def encode_cls(r):
enc = tok(r["text"], truncation=True, max_length=MAX)
enc["labels"] = lid[r["label"]]
return enc
def collate(batch):
n = max(len(b["input_ids"]) for b in batch)
out = {"input_ids": [], "attention_mask": [], "labels": []}
for b in batch:
pad = n - len(b["input_ids"])
out["input_ids"].append(b["input_ids"] + [tok.pad_token_id] * pad)
out["attention_mask"].append(b["attention_mask"] + [0] * pad)
out["labels"].append(b["labels"] + [-100] * pad if spans else b["labels"])
return {k: torch.tensor(v) for k, v in out.items()}
encode = encode_spans if spans else encode_cls
train = [encode(r) for r in train_rows]
Model = AutoModelForTokenClassification if spans else AutoModelForSequenceClassification
model = Model.from_pretrained(BASE, num_labels=len(labels), id2label=dict(enumerate(labels)), label2id=lid).to(dev)
EPOCHS = int(os.environ.get("EPOCHS", 8 if spans else 6)) # EPOCHS=20 for many-class models with few rows each
loader = DataLoader(train, batch_size=16, shuffle=True, collate_fn=collate)
opt = torch.optim.AdamW(model.parameters(), lr=5e-5, weight_decay=0.01)
sched = get_linear_schedule_with_warmup(opt, len(loader) * EPOCHS // 10, len(loader) * EPOCHS)
model.train()
for ep in range(EPOCHS):
total = 0.0
for batch in loader:
with torch.autocast(dev, dtype=torch.float16, enabled=dev == "cuda"):
loss = model(**{k: v.to(dev) for k, v in batch.items()}).loss
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
opt.step(), sched.step(), opt.zero_grad()
total += loss.item()
print(f"epoch {ep + 1}: loss {total / len(loader):.4f}")
model.eval()
def predict_spans(text):
enc = tok(text, truncation=True, max_length=MAX, return_offsets_mapping=True, return_tensors="pt")
offs = enc.pop("offset_mapping")[0].tolist()
with torch.no_grad():
ids = model(**{k: v.to(dev) for k, v in enc.items()}).logits[0].argmax(-1).tolist()
found, cur = [], None
for (s, e), i in zip(offs, ids):
if s == e:
continue
tag = labels[i]
if tag == "O":
cur = None
continue
kind = tag[2:]
if tag.startswith("I-") and cur and cur["label"] == kind:
cur["end"] = e
else:
cur = {"label": kind, "start": s, "end": e}
found.append(cur)
for f in found: # offsets from a prefix-space tokenizer can start on the space
while f["start"] < f["end"] and text[f["start"]] == " ":
f["start"] += 1
return found
if spans:
tp, fp, fn = Counter(), Counter(), Counter()
for r in test_rows:
gold = {(e["label"], e["start"], e["end"]) for e in r["entities"]}
pred = {(e["label"], e["start"], e["end"]) for e in predict_spans(r["text"])}
for g in gold:
(tp if g in pred else fn)[g[0]] += 1
for p in pred - gold:
fp[p[0]] += 1
f1 = lambda t, p, n: (2 * t / (2 * t + p + n)) if t else 0.0
per = {t: {"precision": round(tp[t] / (tp[t] + fp[t]), 3) if tp[t] + fp[t] else 0.0,
"recall": round(tp[t] / (tp[t] + fn[t]), 3) if tp[t] + fn[t] else 0.0,
"f1": round(f1(tp[t], fp[t], fn[t]), 3), "support": tp[t] + fn[t]} for t in types}
T, P, N = sum(tp.values()), sum(fp.values()), sum(fn.values())
scores = {"metric": "exact entity match", "micro_precision": round(T / (T + P), 3), "micro_recall": round(T / (T + N), 3),
"micro_f1": round(f1(T, P, N), 3), "per_type": per}
else:
pred, gold = [], []
for r in test_rows:
enc = tok(r["text"], truncation=True, max_length=MAX, return_tensors="pt")
with torch.no_grad():
pred.append(labels[model(**{k: v.to(dev) for k, v in enc.items()}).logits[0].argmax().item()])
gold.append(r["label"])
acc = sum(p == g for p, g in zip(pred, gold)) / len(gold)
per = {}
for l in labels:
t = sum(p == g == l for p, g in zip(pred, gold))
p_ = sum(p == l for p in pred)
g_ = sum(g == l for g in gold)
per[l] = {"f1": round(2 * t / (p_ + g_), 3) if p_ + g_ else 0.0, "support": g_}
scores = {"metric": "accuracy and macro F1", "accuracy": round(acc, 3),
"macro_f1": round(sum(v["f1"] for v in per.values() if v["support"]) / sum(1 for v in per.values() if v["support"]), 3),
"per_label": per}
scores.update({"base_model": BASE, "train_rows": len(train_rows), "test_rows": len(test_rows), "epochs": EPOCHS, "device": dev})
(folder / "scores.json").write_text(json.dumps(scores, indent=2) + "\n", encoding="utf-8")
print(json.dumps({k: v for k, v in scores.items() if k not in ("per_type", "per_label")}))
model.save_pretrained(folder / "out", safe_serialization=True)
tok.save_pretrained(folder / "out")
|