diff --git a/label_arwiki_hf.py b/label_arwiki_hf.py index d736baf..1784a22 100644 --- a/label_arwiki_hf.py +++ b/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} " diff --git a/train_arabic_r8_hf.py b/train_arabic_r8_hf.py index d390a69..299e095 100644 --- a/train_arabic_r8_hf.py +++ b/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()