| """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)) |
| 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: |
| 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") |
|
|
