Repository Public Research preview
1#!/usr/bin/env python32"""Rewrite-vs-original discriminator D: fine-tune a pretrained encoder (default roberta-large, MIT) to tell Medium *originals*3(label 1, pre-2018 human prose) from their LLM *rewrites* (label 0) on the TRAIN split; model selection on val. Used as (a) a4reranking signal for candidate selection and (b) a reward term for policy optimisation (HUMPA/AuthorMist-style, with KL and5source-fidelity constraints). D is an in-house proxy; the held-out detector (Pangram) is never trained on or tuned against."""6from __future__ import annotations78import argparse9import math10import random11import sys12import time13from pathlib import Path1415import torch1617if torch.cuda.is_available() and hasattr(torch.backends.cuda, "enable_cudnn_sdp"):18 torch.backends.cuda.enable_cudnn_sdp(False) # torch 2.11/cu128: cuDNN SDPA gives NaN gradients in bf16 and rebuilds a plan per shape1920ROOT = Path(__file__).resolve().parents[1]21sys.path.insert(0, str(ROOT / "src"))22from htr.utils.io import append_jsonl, read_jsonl, write_json # noqa: E40223from htr.utils.provenance import run_header # noqa: E402242526def main() -> None:27 ap = argparse.ArgumentParser()28 ap.add_argument("--pairs", required=True)29 ap.add_argument("--splits", required=True)30 ap.add_argument("--run", required=True)31 ap.add_argument("--model-name", default="FacebookAI/roberta-large")32 ap.add_argument("--max-steps", type=int, default=6000)33 ap.add_argument("--batch-size", type=int, default=32)34 ap.add_argument("--lr", type=float, default=1e-5)35 ap.add_argument("--eval-every", type=int, default=500)36 ap.add_argument("--max-len", type=int, default=384)37 ap.add_argument("--limit", type=int, default=None)38 ap.add_argument("--seed", type=int, default=1)39 ap.add_argument("--device", default="cuda")40 a = ap.parse_args()41 from transformers import AutoModelForSequenceClassification, AutoTokenizer4243 torch.manual_seed(a.seed)44 random.seed(a.seed)45 dev = torch.device(a.device)46 tok = AutoTokenizer.from_pretrained(a.model_name)47 model = AutoModelForSequenceClassification.from_pretrained(a.model_name, num_labels=2).to(dev)48 splits = {s["doc_id"]: s["split"] for s in read_jsonl(a.splits)}49 pairs = list(read_jsonl(a.pairs))50 # one original per passage (originals repeat across the 6 rewrites of a passage) + all rewrites → balanced by sampling51 def build(split):52 seen, rows = set(), []53 for p in pairs:54 if splits.get(p["doc_id"]) != split:55 continue56 rows.append((p["source_text"], 0))57 pid = p.get("passage_id", p["pair_id"])58 if pid not in seen:59 seen.add(pid)60 rows.append((p["target_text"], 1))61 return rows62 tr, va = build("train"), build("val")63 if a.limit:64 random.shuffle(tr); tr = tr[: a.limit]65 va = va[:4000]66 pos = sum(l for _, l in tr); print(f"train rows {len(tr)} (originals {pos}), val rows {len(va)}", flush=True)67 w = torch.tensor([pos / max(1, len(tr) - pos), 1.0], device=dev) # class weights balance rewrites (6 per passage) vs originals68 run_dir = ROOT / "artifacts/runs" / a.run; run_dir.mkdir(parents=True, exist_ok=True)69 opt = torch.optim.AdamW(model.parameters(), lr=a.lr, weight_decay=0.01)70 ac = torch.autocast("cuda", dtype=torch.bfloat16) if dev.type == "cuda" else None7172 def batches(rows, shuffle):73 idx = list(range(len(rows)))74 if shuffle:75 random.shuffle(idx)76 for i in range(0, len(idx), a.batch_size):77 ch = [rows[j] for j in idx[i : i + a.batch_size]]78 enc = tok([t for t, _ in ch], return_tensors="pt", padding=True, truncation=True, max_length=a.max_len, pad_to_multiple_of=32)79 yield {k: v.to(dev) for k, v in enc.items()}, torch.tensor([l for _, l in ch], device=dev)8081 @torch.no_grad()82 def evaluate():83 model.eval(); scores, labels = [], []84 for enc, y in batches(va, False):85 with (ac or torch.no_grad()):86 lg = model(**enc).logits.float()87 scores += lg.softmax(-1)[:, 1].tolist(); labels += y.tolist()88 model.train()89 import numpy as np90 s, l = np.array(scores), np.array(labels)91 acc = float(((s > 0.5) == (l == 1)).mean())92 # AUROC by rank statistic93 order = s.argsort(); ranks = np.empty(len(s)); ranks[order] = np.arange(1, len(s) + 1)94 n1, n0 = int(l.sum()), int((l == 0).sum())95 auroc = float((ranks[l == 1].sum() - n1 * (n1 + 1) / 2) / max(1, n1 * n0))96 return {"val_acc": acc, "val_auroc": auroc, "mean_p_human_orig": float(s[l == 1].mean()), "mean_p_human_rewrite": float(s[l == 0].mean())}9798 step, best, t0 = 0, -1.0, time.time()99 model.train()100 while step < a.max_steps:101 for enc, y in batches(tr, True):102 for g in opt.param_groups:103 g["lr"] = a.lr * min(1.0, (step + 1) / 300) * 0.5 * (1 + math.cos(math.pi * min(1.0, step / a.max_steps)))104 with (ac or torch.no_grad()):105 lg = model(**enc).logits.float()106 loss = torch.nn.functional.cross_entropy(lg, y, weight=w)107 loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0); opt.step(); opt.zero_grad(set_to_none=True)108 step += 1109 if step % 25 == 0:110 append_jsonl(run_dir / "train_log.jsonl", {"step": step, "loss": float(loss.detach()), "elapsed_s": round(time.time() - t0, 1)})111 if step % a.eval_every == 0 or step >= a.max_steps:112 ev = evaluate(); append_jsonl(run_dir / "train_log.jsonl", {"step": step, "eval": ev}); print(step, ev, flush=True)113 if ev["val_auroc"] > best:114 best = ev["val_auroc"]; model.save_pretrained(run_dir / "best"); tok.save_pretrained(run_dir / "best")115 if step >= a.max_steps:116 break117 write_json(run_dir / "train_summary.json", {**run_header({"run": a.run}), "model_name": a.model_name, "steps": step, "best_val_auroc": best, "n_train_rows": len(tr), "n_val_rows": len(va), "wall_seconds": round(time.time() - t0, 1), "args": vars(a)})118 print("best val AUROC", best)119120121if __name__ == "__main__":122 main()123