Add source/train.py

source/train.py

@@ -0,0 +1,160 @@
+"""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")