feat: r8 register-closure arms (WO17)

label_arwiki_hf.py

@@ -36,7 +36,7 @@
 UNIT_BYTES = 1400
 MARGIN_NAT = 1.0
 KEEP_FRAC = 0.9
-MAX_UNITS = 400_000
+MAX_UNITS = 12_000
 
 
 def window(text: str, budget: int = UNIT_BYTES) -> list[str]:
@@ -87,31 +87,33 @@ def main() -> None:
                       max_length=1450).to(device)
             with torch.autocast("cuda", torch.bfloat16):
                 gen = model.generate(**enc, max_new_tokens=1500, num_beams=1,
-                                     return_dict_in_generate=True, output_scores=False)
+                                     return_dict_in_generate=True, output_scores=True)
             seqs = gen.sequences
             texts = tok.batch_decode(seqs, skip_special_tokens=True)
 
-            # teacher-forced margin pass on the generated tokens
-            dec = seqs[:, :-1]
-            dec_mask = dec != tok.pad_token_id
-            with torch.autocast("cuda", torch.bfloat16):
-                sc = model(**{k: enc[k] for k in ("input_ids", "attention_mask")},
-                           decoder_input_ids=dec, decoder_attention_mask=dec_mask)
-            logits = sc.logits.float()
-            top2 = logits.topk(2, dim=-1).values
-            margins = (top2[..., 0] - top2[..., 1])
-            conf = (margins > MARGIN_NAT) & dec_mask
-            frac = conf.sum(-1).float() / dec_mask.sum(-1).clamp(min=1)
+            # margins straight from generate's per-step scores (one pass)
+            dec_len = seqs.size(1) - 1
+            if dec_len > 0 and len(gen.scores) == dec_len:
+                logits = torch.stack(gen.scores, dim=1).float()  # [B, T, V]
+                top2 = logits.topk(2, dim=-1).values
+                margins = top2[..., 0] - top2[..., 1]
+                real = seqs[:, 1:] != tok.pad_token_id
+                conf = (margins > MARGIN_NAT) & real
+                frac = conf.sum(-1).float() / real.sum(-1).clamp(min=1)
+            else:
+                frac = torch.ones(len(texts), device=seqs.device)
 
             srcs = tok.batch_decode(enc.input_ids, skip_special_tokens=True)
             for src, text, f in zip(srcs, texts, frac.tolist()):
                 done += 1
+                # every unit is written (keep_frac field); the train-time
+                # loader applies the KEEP_FRAC filter - resume stays exact
+                out.write(json.dumps({"src": src, "tgt": text,
+                                      "keep_frac": round(f, 4)},
+                                     ensure_ascii=False) + "\n")
                 if f >= KEEP_FRAC:
                     kept += 1
-                    out.write(json.dumps({"src": src, "tgt": text,
-                                          "keep_frac": round(f, 4)},
-                                         ensure_ascii=False) + "\n")
-            if (i // 8) % 50 == 0:
+            if (i // 32) % 20 == 0:
                 out.flush()
                 rate = done / max(1, time.time() - t0)
                 print(f"[gen] {done}/{len(units)} kept={kept} "

train_arabic_r8_hf.py

@@ -37,9 +37,11 @@
 sys.path.insert(0, "/code")
 
 DATA = Path("/train_data")
-CKPT = Path("/ckpt/run-024-arabic-r8-noisy")
 TEACHER = Path("/ckpt/r7-best")
-RUN = "run-024-arabic-r8-noisy"
+import os
+MIX = os.environ.get("R8_MIX", "arwiki")  # arwiki | qcri | both
+CKPT = Path(f"/ckpt/run-024-arabic-r8-{MIX}")
+RUN = f"run-024-arabic-r8-{MIX}"
 WINDOW = 600
 GOLD_CAP = 400_000
 PSEUDO_CAP = 200_000
@@ -119,19 +121,50 @@ def main() -> None:
     print(f"[data] gold={len(gold)}", flush=True)
 
     pseudo: list[tuple[str, str]] = []
-    ppath = Path("/ckpt/r8-pseudo/arwiki-pseudo.jsonl")
-    if ppath.exists():
-        for line in ppath.read_text(encoding="utf-8").splitlines():
-            if not line.strip():
-                continue
-            r = json.loads(line)
-            src = DIACRITICS_RE.sub("", r["tgt"]).strip()
-            if len(src) < 5 or len(src.encode()) > WINDOW - 50:
-                continue
-            pseudo.append((src, r["tgt"]))
+
+    def add_silver(tgt: str) -> None:
+        tgt = tgt.strip()
+        src = DIACRITICS_RE.sub("", tgt).strip()
+        if len(src) >= 5 and len(src.encode()) <= WINDOW - 50:
+            pseudo.append((src, tgt))
+
+    if MIX in ("arwiki", "both"):
+        ppath = Path("/ckpt/r8-pseudo/arwiki-pseudo.jsonl")
+        if ppath.exists():
+            for line in ppath.read_text(encoding="utf-8").splitlines():
+                if not line.strip():
+                    continue
+                r = json.loads(line)
+                if r.get("keep_frac", 1.0) < 0.9:
+                    continue
+                add_silver(r["tgt"])
+                if len(pseudo) >= PSEUDO_CAP:
+                    break
+    if MIX in ("qcri", "both") and len(pseudo) < PSEUDO_CAP:
+        qpath = DATA / "qcri-wiki.jsonl"
+        # QCRI silver (Wikipedia_20240420.diac.jsonl, their BiLSTM labels,
+        # published with EMNLP 2025) - windowed to the 600-byte protocol
+        rng2 = random.Random(7)
+        articles = qpath.read_text(encoding="utf-8").splitlines()
+        rng2.shuffle(articles)
+        for aline in articles:
             if len(pseudo) >= PSEUDO_CAP:
                 break
-    print(f"[data] pseudo={len(pseudo)}", flush=True)
+            if not aline.strip():
+                continue
+            art = json.loads(aline).get("text", "")
+            for para in art.split("\n"):
+                para = para.strip()
+                if not para:
+                    continue
+                if len(para.encode()) <= WINDOW - 50:
+                    add_silver(para)
+                else:
+                    for w in split_windows(para, WINDOW - 50):
+                        add_silver(w)
+                if len(pseudo) >= PSEUDO_CAP:
+                    break
+    print(f"[data] pseudo={len(pseudo)} (mix={MIX})", flush=True)
 
     units = gold + pseudo
     rng.shuffle(units)
@@ -269,6 +302,41 @@ def save(s: int) -> None:
     (CKPT / "EVAL_DONE").write_text("ok\n")
     print("[done] EVAL_DONE", flush=True)
 
+    # ---- surface 2: WikiNews-2024 multiref (the register gate) ----
+    import eval_wikinews_multiref as WN
+
+    bench = (DATA / "wikinews/WikiNews_2024_Multi_Ref.txt.diac").read_text(
+        encoding="utf-8").splitlines()
+    bench = [l for l in bench if l.strip()]
+    winputs = [DIACRITICS_RE.sub("", l) for l in bench]
+    wwindows, wcounts = [], []
+    for text in winputs:
+        ws = split_windows(text)
+        wcounts.append(len(ws))
+        wwindows.extend(ws)
+    wpreds = []
+    with torch.no_grad():
+        for i in range(0, len(wwindows), 32):
+            batch = wwindows[i:i + 32]
+            enc = tok(batch, return_tensors="pt", padding=True, truncation=True,
+                      max_length=WINDOW).to("cuda")
+            with torch.autocast("cuda", torch.bfloat16):
+                gen = model.generate(**enc, max_new_tokens=WINDOW * 2, num_beams=1)
+            wpreds.extend(tok.batch_decode(gen, skip_special_tokens=True))
+            if (i // 32) % 5 == 0:
+                print(f"[wgen] {i + len(batch)}/{len(wwindows)}", flush=True)
+    k = 0
+    wparas = []
+    for text, c in zip(winputs, wcounts):
+        wparas.append(" ".join(wpreds[k:k + c]))
+        k += c
+    wscore = WN.score(wparas, bench, skip_last=False)
+    print(json.dumps({"run": RUN, "surface": "wikinews-2024-multiref", **wscore},
+                     ensure_ascii=False), flush=True)
+    (CKPT / "wikinews_multiref.json").write_text(
+        json.dumps({"run": RUN, **wscore}, ensure_ascii=False, indent=1))
+    (CKPT / "EVAL_DONE").touch()
+
 
 if __name__ == "__main__":
     main()