From fc65826a104c4464d8d55bd39da06c8f37afda08 Mon Sep 17 00:00:00 2001 From: Ronald Tse Date: Fri, 30 Oct 2026 15:47:05 +0800 Subject: [PATCH 1/5] feat(WO64): B3 Option-B oracle relabel via shipped ONNX arms + fresh-student stage 2 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Stage 1 decodes the two surviving SHIPPED artifacts (ara-diac-2.0 r7 ByT5-base kv; ara-diac-news-1.0 int8 plain) with length-bucketed zero-pad batches — parity-tested against the house runtime: kv batch matches exactly, plain exports are batch=1-traced and decode single-row. Stage 2 trains a fresh google/byt5-base student on the oracle targets with the gate re-baselined to the shipped arms' published numbers. --- oracle_relabel_onnx.py | 399 +++++++++++++++++++++++++++++++++++++ test_oracle_onnx_parity.py | 71 +++++++ train_oracle_student_hf.py | 253 +++++++++++++++++++++++ 3 files changed, 723 insertions(+) create mode 100644 oracle_relabel_onnx.py create mode 100644 test_oracle_onnx_parity.py create mode 100644 train_oracle_student_hf.py diff --git a/oracle_relabel_onnx.py b/oracle_relabel_onnx.py new file mode 100644 index 0000000..495f3ea --- /dev/null +++ b/oracle_relabel_onnx.py @@ -0,0 +1,399 @@ +"""WO64/B3 Option-B stage 1 — oracle relabel via the two surviving SHIPPED arms. + +The WO39 script (distill_oracle_hf.py) decodes PyTorch arm checkpoints that +the storage layer no longer holds (the WO64 verdict); the Option-B arms are +the shipped IMF artifacts themselves: + + r7 = /ckpt/b3-arms/ara-diac-2.0-fp32.zip ByT5-base fp32, kv decoder + news = /ckpt/b3-arms/ara-diac-news-1.0.zip ByT5-small int8, plain decoder + +Decoding is ONNX with LENGTH-BUCKETED batches: every row in a batch has an +identical encoded length, so the encoder never sees padding and the WO51 +unmasked-encoder hazard cannot apply. Rows follow run-047's corpus recipe +exactly (article shuffle seed 7, row shuffle seed 42, WINDOW=600, +morpheme-slash strip); the per-row oracle is the arm prediction with minimum +char edit distance to the corpus silver, ties -> r7. + +Env: D_RUN (default run-084-oracle-relabel-b), D_SUBSET (default 300000), +D_SMOKE (set -> run ~100 rows end to end and exit), D_BS (max batch, +default 48). + +Durability: per-arm raw preds are appended keyed by row index and +force-synced to the bucket every ~4k decoded rows; on relaunch each arm +resumes from its synced raw file. Consolidation runs only when both arms +cover all rows. + +Usage (a100-large): + python /job/oracle_relabel_onnx.py +""" + +from __future__ import annotations + +import hashlib +import json +import os +import random +import re +import time +import zipfile +from pathlib import Path + +import numpy as np + +DATA = Path("/train_data") +CKPT = Path("/ckpt") / os.environ.get("D_RUN", "run-084-oracle-relabel-b") +WINDOW = 600 +PSEUDO_CAP = 900_000 +SUBSET = int(os.environ.get("D_SUBSET", "300000")) +SMOKE = bool(os.environ.get("D_SMOKE")) +BS = int(os.environ.get("D_BS", "48")) +ARMS = [ + ("r7", "/ckpt/b3-arms/ara-diac-2.0-fp32.zip"), + ("news", "/ckpt/b3-arms/ara-diac-news-1.0.zip"), +] +DIACRITICS_RE = re.compile("[ؐ-ًؚ-ٰٟۖ-ۜ۟-۪ۨ-ۭ]") + +BYTE_OFFSET = 3 +PAD_ID = 0 +EOS_ID = 1 + + +def encode(text: str) -> list[int]: + return [b + BYTE_OFFSET for b in text.encode("utf-8")] + [EOS_ID] + + +def decode_ids(ids: list[int]) -> str: + out = bytearray() + for t in ids: + if t == EOS_ID: + break + if t in (PAD_ID, 2): + continue + out.append((t - BYTE_OFFSET) % 256) + return out.decode("utf-8", errors="replace") + + +def lev(a: str, b: str) -> int: + try: + from rapidfuzz.distance import Levenshtein + + return Levenshtein.distance(a, b) + except ImportError: + prev = list(range(len(b) + 1)) + for i, ca in enumerate(a, 1): + cur = [i] + for j, cb in enumerate(b, 1): + cur.append(min(prev[j] + 1, cur[j - 1] + 1, + prev[j - 1] + (ca != cb))) + prev = cur + return prev[-1] + + +class Arm: + def __init__(self, name: str, zip_path: str): + import onnxruntime as ort + + self.name = name + opts = ort.SessionOptions() + opts.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL + graphs, digests = {}, {} + with zipfile.ZipFile(zip_path) as z: + for member in z.namelist(): + if member.endswith(".onnx"): + graphs[member] = z.read(member) + elif member == "metadata.yaml": + meta = z.read(member).decode("utf-8") + in_sha = False + for line in meta.splitlines(): + if line.strip() == "sha256:": + in_sha = True + continue + if in_sha: + m = re.match(r"\s+([\w.-]+): ([0-9a-f]{64})", line) + if m: + digests[m.group(1)] = m.group(2) + elif line.strip() and not line.startswith(" "): + in_sha = False + for fname, blob in graphs.items(): + want = digests.get(fname) + if want: + got = hashlib.sha256(blob).hexdigest() + if got != want: + raise RuntimeError(f"{zip_path}:{fname} sha256 mismatch") + providers = [p for p in ("CUDAExecutionProvider", "CPUExecutionProvider") + if p in ort.get_available_providers()] + self.enc = ort.InferenceSession(graphs["encoder.onnx"], opts, providers) + dname = ("decoder-kv.onnx" + if "decoder-kv.onnx" in graphs else "decoder.onnx") + self.kv = "decoder-kv.onnx" in graphs + self.dec = ort.InferenceSession(graphs[dname], opts, providers) + self.out_names = [o.name for o in self.dec.get_outputs()] + self.pasts = {} + for m in self.dec.get_inputs(): + if m.name.startswith("past_"): + shape = m.shape # [batch, heads, past_seq, d_kv] + heads = shape[1] if isinstance(shape[1], int) else 4 + d_kv = shape[3] if isinstance(shape[3], int) else 8 + self.pasts[m.name] = np.zeros((1, heads, 0, d_kv), np.float32) + + def decode_equal(self, srcs: list[str]) -> list[str]: + """Greedy decode a batch of identical-encoded-length rows.""" + ids = [encode(s) for s in srcs] + n = len(srcs) + enc = np.array(ids, dtype=np.int64) + hidden = self.enc.run(None, {"input_ids": enc})[0] + max_len = min(1300, 2 * enc.shape[1] + 8) + outs: list[list[int]] = [[] for _ in range(n)] + alive = np.ones(n, dtype=bool) + if self.kv: + pasts = {k: np.repeat(z, n, axis=0) for k, z in self.pasts.items()} + current = np.full((n, 1), PAD_ID, dtype=np.int64) + for _ in range(max_len): + res = dict(zip(self.out_names, self.dec.run( + None, {"input_ids": current, + "encoder_hidden_states": hidden, **pasts}))) + tok = np.argmax(res["logits"][:, -1, :], axis=1) + for b in np.nonzero(alive)[0]: + t = int(tok[b]) + if t == EOS_ID: + alive[b] = False + else: + outs[b].append(t) + if not alive.any(): + break + pasts = {name: res[name.replace("past_", "present_", 1)] + for name in pasts} + current = tok.reshape(n, 1) + else: + current = np.full((n, 1), PAD_ID, dtype=np.int64) + for _ in range(max_len): + logits = self.dec.run( + None, {"input_ids": current, + "encoder_hidden_states": hidden})[0] + tok = np.argmax(logits[:, -1, :], axis=1) + for b in np.nonzero(alive)[0]: + t = int(tok[b]) + if t == EOS_ID: + alive[b] = False + else: + outs[b].append(t) + if not alive.any(): + break + current = np.concatenate([current, tok.reshape(n, 1)], axis=1) + return [decode_ids(o) for o in outs] + + +def split_windows(text: str, budget: int = WINDOW) -> list[str]: + if len(text.encode("utf-8")) <= budget: + return [text] + words, cur, n, wins = text.split(), [], 0, [] + for w in words: + c = len(w.encode("utf-8")) + 1 + if cur and n + c > budget: + wins.append(" ".join(cur)) + cur, n = [], 0 + cur.append(w) + n += c + if cur: + wins.append(" ".join(cur)) + return wins + + +def build_rows() -> list[tuple[str, str]]: + """run-047's exact qcri-only silver corpus, stripped, shuffled.""" + qpath = DATA / "qcri-wiki.jsonl" + articles = qpath.read_text(encoding="utf-8").splitlines() + rng2 = random.Random(7) + rng2.shuffle(articles) + pseudo: list[tuple[str, str]] = [] + + 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)) + + for aline in articles: + if len(pseudo) >= PSEUDO_CAP: + break + 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 + stripped = [] + for _src, tgt in pseudo: + tgt2 = " ".join(t.split("/", 1)[0] for t in tgt.split()) + src2 = DIACRITICS_RE.sub("", tgt2).strip() + if len(src2) >= 5: + stripped.append((src2, tgt2)) + rng = random.Random(42) + rng.shuffle(stripped) + rows = stripped[:SUBSET] + print(f"[data] stripped={len(stripped)} subset={len(rows)}", flush=True) + return rows + + +def sync_to_bucket(path: Path) -> None: + from huggingface_hub import HfFileSystem + + fs = HfFileSystem(token=os.environ.get("HF_TOKEN")) + dst = f"buckets/Interscript/isx-training/{path.relative_to(Path('/ckpt'))}" + with open(path, "rb") as f, fs.open(dst, "wb") as g: + g.write(f.read()) + if not fs.exists(dst): + raise RuntimeError(f"sync verify failed: {dst}") + + +def decode_groups(nbytes: list[int], done: set[int], + bs: int = BS) -> list[list[int]]: + """Deterministic length-grouped decode plan over remaining rows.""" + by_len: dict[int, list[int]] = {} + for i, nb in enumerate(nbytes): + if i not in done: + by_len.setdefault(nb, []).append(i) + groups = [] + for nb in sorted(by_len): + chunk = by_len[nb] + for j in range(0, len(chunk), bs): + groups.append(chunk[j:j + bs]) + sizes = [len(g) for g in groups] + print(f"[plan] groups={len(groups)} rows={sum(sizes)} " + f"bs={bs} mean={sum(sizes) / max(1, len(sizes)):.1f}", flush=True) + return groups + + +def decode_arm(arm: Arm, srcs: list[str]) -> None: + raw = CKPT / f"preds-{arm.name}.raw.jsonl" + done: set[int] = set() + if raw.exists(): + for line in raw.read_text(encoding="utf-8").splitlines(): + if line.strip(): + done.add(json.loads(line)["i"]) + print(f"[{arm.name}] resume: {len(done)} rows already decoded", + flush=True) + nbytes = [len(s.encode()) + 1 for s in srcs] + # plain-decoder exports are batch=1-traced (parity test: batched + # plain decode corrupts; only the kv graph is batch-safe) — so the + # int8 news arm always decodes single-row, exactly like the runtime + groups = decode_groups(nbytes, done, BS if arm.kv else 1) + t0 = time.time() + n_done = 0 + with raw.open("a", encoding="utf-8") as f: + for gi, g in enumerate(groups): + preds = arm.decode_equal([srcs[i] for i in g]) + for i, p in zip(g, preds): + f.write(json.dumps({"i": i, "p": p}, ensure_ascii=False) + "\n") + n_done += len(g) + f.flush() + if gi % 40 == 0: + try: + sync_to_bucket(raw) + except Exception as e: # noqa: BLE001 + print(f"[{arm.name}] sync failed: {e}", flush=True) + if gi % 20 == 0: + r = n_done / max(1, time.time() - t0) + rem = len(srcs) - len(done) - n_done + print(f"[{arm.name}] {n_done}/{len(srcs) - len(done)} " + f"({r:.0f} rows/s, eta {rem / max(r, 1) / 3600:.1f}h)", + flush=True) + try: + sync_to_bucket(raw) + except Exception as e: # noqa: BLE001 + print(f"[{arm.name}] final sync failed: {e}", flush=True) + + +def raw_preds(name: str, n: int) -> list[str] | None: + raw = CKPT / f"preds-{name}.raw.jsonl" + if not raw.exists(): + return None + by_i: dict[int, str] = {} + for line in raw.read_text(encoding="utf-8").splitlines(): + if line.strip(): + r = json.loads(line) + by_i[r["i"]] = r["p"] + if len(by_i) < n: + print(f"[{name}] raw covers {len(by_i)}/{n} rows", flush=True) + return None + return [by_i[i] for i in range(n)] + + +def consolidate(rows: list[tuple[str, str]], + preds: dict[str, list[str]]) -> None: + with (CKPT / "arm-preds.jsonl").open("w", encoding="utf-8") as f, \ + (CKPT / "oracle-targets.jsonl").open("w", encoding="utf-8") as g: + cers = {n: 0.0 for n, _ in ARMS} + oracle_cer = 0.0 + disagree = 0 + for i, (src, silver) in enumerate(rows): + row_preds = {n: preds[n][i] for n, _ in ARMS} + if len(set(row_preds.values())) > 1: + disagree += 1 + best_name, best_d = ARMS[0][0], lev(row_preds[ARMS[0][0]], silver) + best_tgt = row_preds[ARMS[0][0]] + for n, _ in ARMS[1:]: + d = lev(row_preds[n], silver) + if d < best_d: + best_name, best_d, best_tgt = n, d, row_preds[n] + for n in cers: + cers[n] += lev(row_preds[n], silver) / max(1, len(silver)) + oracle_cer += best_d / max(1, len(silver)) + f.write(json.dumps({"src": src, "silver": silver, **row_preds}, + ensure_ascii=False) + "\n") + g.write(json.dumps({"src": src, "tgt": best_tgt, + "pick": best_name}, ensure_ascii=False) + "\n") + n = len(rows) + picks: dict[str, int] = {} + for l in (CKPT / "oracle-targets.jsonl").read_text( + encoding="utf-8").splitlines(): + p = json.loads(l)["pick"] + picks[p] = picks.get(p, 0) + 1 + stats = {"rows": n, "disagreement": round(disagree / n, 4), + "arm_cer": {k: round(v / n, 4) for k, v in cers.items()}, + "oracle_cer": round(oracle_cer / n, 4), "picks": picks} + print(json.dumps(stats), flush=True) + (CKPT / "relabel_stats.json").write_text(json.dumps(stats, indent=1)) + + +def main() -> None: + CKPT.mkdir(parents=True, exist_ok=True) + rows = build_rows() + srcs = [s for s, _ in rows] + loaded = [(n, Arm(n, p)) for n, p in ARMS] + for name, arm in loaded: + decode_arm(arm, srcs) + preds = {} + for name, _ in loaded: + p = raw_preds(name, len(rows)) + if p is None: + print(f"[consolidate] waiting on arm {name}", flush=True) + return + preds[name] = p + consolidate(rows, preds) + (CKPT / "RELABEL_DONE").write_text("ok\n") + try: + sync_to_bucket(CKPT / "RELABEL_DONE") + except Exception as e: # noqa: BLE001 + print(f"sync failed: {e}", flush=True) + + +if __name__ == "__main__": + if SMOKE: + os.environ["D_SUBSET"] = "100" + SUBSET = 100 + main() + if SMOKE: + (CKPT / "SMOKE_DONE").write_text("ok\n") + try: + sync_to_bucket(CKPT / "SMOKE_DONE") + except Exception: + pass diff --git a/test_oracle_onnx_parity.py b/test_oracle_onnx_parity.py new file mode 100644 index 0000000..a3c8b64 --- /dev/null +++ b/test_oracle_onnx_parity.py @@ -0,0 +1,71 @@ +"""Parity: oracle_relabel_onnx.Arm.decode_equal vs house runtime translate. + +Three equal-byte-length Arabic variants (1:1 letter substitutions preserve +byte length) decoded greedy through both paths must produce identical +strings. Run on CPU against the shipped zips in /tmp/b3-restage. +""" + +import sys +import time +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).parent)) +sys.path.insert(0, "/Users/mulgogi/src/interscript/interscript-models/runtime/src") + +from interscript_ml import Model # noqa: E402 +from oracle_relabel_onnx import Arm # noqa: E402 + +T1 = "العلم نور والجهل ظلام والكتاب خير جليس" +T2 = T1.translate(str.maketrans("لم", "مو")) +T3 = T1.translate(str.maketrans("جهل", "زحم")) +TEXTS = [T1, T2, T3] + +lens = {len(t.encode()) for t in TEXTS} +assert len(lens) == 1, f"byte lengths differ: {lens}" +L = lens.pop() + 1 # + trailing EOS +MAX_LEN = min(1300, 2 * L + 8) +print(f"[parity] texts={len(TEXTS)} bytes={L - 1} max_len={MAX_LEN}") + +ZIP = "/tmp/b3-restage/ara-diac-news-1.0.zip" +t0 = time.time() +house = Model.load(ZIP) +ref = [house.translate(t, max_len=MAX_LEN) for t in TEXTS] +print(f"[parity] house runtime done in {time.time() - t0:.1f}s") + +arm = Arm("news", ZIP) +t0 = time.time() +# plain-decoder exports are batch=1-traced: the relabel plan decodes +# them single-row; parity must hold in that mode +mine = [arm.decode_equal([t])[0] for t in TEXTS] +print(f"[parity] single-row decode done in {time.time() - t0:.1f}s") + +ok = True +for i, (r, m) in enumerate(zip(ref, mine)): + match = r == m + ok &= match + print(f"[parity] row {i}: {'MATCH' if match else 'MISMATCH'}") + if not match: + print(f" house: {r!r}") + print(f" mine: {m!r}") + +ZIP = "/tmp/b3-restage/ara-diac-2.0-fp32.zip" +t0 = time.time() +house = Model.load(ZIP) +ref = [house.translate(t, max_len=MAX_LEN) for t in TEXTS] +print(f"[parity] house r7 done in {time.time() - t0:.1f}s") + +arm = Arm("r7", ZIP) +t0 = time.time() +mine = arm.decode_equal(TEXTS) +print(f"[parity] batched r7 decode done in {time.time() - t0:.1f}s") + +for i, (r, m) in enumerate(zip(ref, mine)): + match = r == m + ok &= match + print(f"[parity] r7 row {i}: {'MATCH' if match else 'MISMATCH'}") + if not match: + print(f" house: {r!r}") + print(f" mine: {m!r}") + +print("[parity] PASS" if ok else "[parity] FAIL") +sys.exit(0 if ok else 1) diff --git a/train_oracle_student_hf.py b/train_oracle_student_hf.py new file mode 100644 index 0000000..aa337b7 --- /dev/null +++ b/train_oracle_student_hf.py @@ -0,0 +1,253 @@ +"""WO64/B3 Option-B stage 2 — train a fresh byt5-base student on oracle targets. + +The WO39 continue-train needs the run-047/best PyTorch init, which the +storage layer no longer holds (WO64 verdict); this derivative starts from +stock google/byt5-base on the same oracle-targets.jsonl (src, tgt, pick), +same lr schedule, then the same eval surfaces in-job (SadeedDiac-25 +default protocol + WikiNews-2024 multiref RAW and first-alt, project_haraqat +on stitched windows). + +Gate re-baselined on the shipped arms (the stable references): +- RAW multiref WER < 17.3794 (shipped r7, ara-diac-2.0) — must beat r7; +- stretch RAW multiref WER <= 10.1262 (shipped news arm, ara-diac-news-1.0); +- SadeedDiac default DER within +0.1 of 2.2864 (shipped r7) or better. + +Env: O_RUN (default run-084-oracle-relabel-b), O_EPOCHS (default 3). + +Usage (a100-large, mounts arabic-r8 dataset + isx-training bucket): + python /job/train_oracle_student_hf.py +""" + +from __future__ import annotations + +import json +import os +import random +from pathlib import Path + +import torch +from torch.utils.data import DataLoader, Dataset + +DATA = Path("/train_data") +CKPT = Path("/ckpt") / os.environ.get("O_RUN", "run-084-oracle-relabel-b") +ORACLE = CKPT / "oracle-targets.jsonl" +INIT = "google/byt5-base" +EPOCHS = int(os.environ.get("O_EPOCHS", "3")) +WINDOW = 600 +BS = 8 +ACCUM = 2 + +DIACRITICS_RE = __import__("re").compile("[ؐ-ًؚ-ٰٟۖ-ۜ۟-۪ۨ-ۭ]") + + +def split_windows(text: str, budget: int = WINDOW) -> list[str]: + if len(text.encode("utf-8")) <= budget: + return [text] + words, cur, n, wins = text.split(), [], 0, [] + for w in words: + c = len(w.encode("utf-8")) + 1 + if cur and n + c > budget: + wins.append(" ".join(cur)); cur, n = [], 0 + cur.append(w); n += c + if cur: + wins.append(" ".join(cur)) + return wins + + +def project_haraqat(pred: str, text: str) -> str: + from difflib import SequenceMatcher + + pred_haraqat = [""] + for ch in pred: + if DIACRITICS_RE.match(ch): + pred_haraqat[-1] += ch + else: + pred_haraqat.append("") + pred_haraqat = pred_haraqat[1:] + pred_letters = [c for c in pred if not DIACRITICS_RE.match(c)] + text_letters = [c for c in text if not DIACRITICS_RE.match(c)] + sm = SequenceMatcher(None, text_letters, pred_letters, autojunk=False) + out = [] + for op, i1, i2, j1, j2 in sm.get_opcodes(): + if op == "equal": + for k in range(i2 - i1): + out.append(text_letters[i1 + k] + pred_haraqat[j1 + k]) + else: + for k in range(i1, i2): + out.append(text_letters[k]) + return "".join(out) + + +def main() -> None: + import sys + + sys.path.insert(0, "/job") + from transformers import (AutoModelForSeq2SeqLM, AutoTokenizer, + get_cosine_schedule_with_warmup) + + units = [(r["src"], r["tgt"]) for r in + (json.loads(l) for l in ORACLE.read_text( + encoding="utf-8").splitlines() if l.strip())] + random.Random(43).shuffle(units) + print(f"[data] oracle units={len(units)}", flush=True) + + tok = AutoTokenizer.from_pretrained(INIT) + model = AutoModelForSeq2SeqLM.from_pretrained(INIT).cuda().train() + model.gradient_checkpointing_enable() + model.config.use_cache = False + + eps = sorted(CKPT.glob("oepoch-*"), key=lambda p: int(p.name.split("-")[1])) + if eps: + model = AutoModelForSeq2SeqLM.from_pretrained(eps[-1]).cuda().train() + model.gradient_checkpointing_enable() + model.config.use_cache = False + print(f"[resume] {eps[-1].name}", flush=True) + + class DS(Dataset): + def __len__(self): + return len(units) + + def __getitem__(self, i): + return units[i] + + def collate(b): + es = [tok(s, truncation=True, max_length=WINDOW).input_ids + for s, _ in b] + ds_ = [tok(t, truncation=True, max_length=WINDOW * 2).input_ids + for _, t in b] + mx, md = max(len(e) for e in es), max(len(d) for d in ds_) + ids = torch.zeros((len(b), mx), dtype=torch.long) + am = torch.zeros((len(b), mx), dtype=torch.long) + lb = torch.full((len(b), md), -100, dtype=torch.long) + for i, (e, d) in enumerate(zip(es, ds_)): + ids[i, :len(e)] = torch.tensor(e) + am[i, :len(e)] = 1 + lb[i, :len(d)] = torch.tensor(d) + return ids, am, lb + + loader = DataLoader(DS(), batch_size=BS, shuffle=True, + collate_fn=collate, num_workers=2, drop_last=True) + total_steps = len(loader) * EPOCHS + opt = torch.optim.AdamW(model.parameters(), lr=3e-5, weight_decay=0.01) + sched = get_cosine_schedule_with_warmup(opt, 200, total_steps) + + step, micro = 0, 0 + for ep in range(EPOCHS): + for ids, am, lb in loader: + ids, am, lb = ids.cuda(), am.cuda(), lb.cuda() + with torch.autocast("cuda", torch.bfloat16): + out = model(input_ids=ids, attention_mask=am, labels=lb) + if not torch.isfinite(out.loss): + print(f"[nan] step {step}", flush=True) + opt.zero_grad() + continue + (out.loss / ACCUM).backward() + micro += 1 + if micro % ACCUM: + continue + torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) + opt.step() + sched.step() + opt.zero_grad() + step += 1 + if step % 100 == 0: + print(f"[train] {step}/{total_steps // ACCUM} " + f"loss={float(out.loss):.3f}", flush=True) + ep_dir = CKPT / f"oepoch-{ep}" + ep_dir.mkdir(exist_ok=True) + model.save_pretrained(ep_dir) + tok.save_pretrained(ep_dir) + print(f"[ckpt] oepoch-{ep} saved", flush=True) + + # ---- eval surfaces (run-047 protocol) ---- + # NOTE: everything below the checkpoint saves is an epilogue — a + # missing dependency here must never fail a completed training run + # (run-080-r8-stage2 died here after 12.7h on `import pandas`). + model.config.use_cache = True + model.eval() + + try: + import pandas as pd # noqa: F401 + import pyarrow.parquet as pq # noqa: F401 + except ImportError as e: + print(f"[epilogue] skipped: {e} — checkpoints are complete; " + "run eval_r8_stage2.py separately", flush=True) + return + + table = pq.read_table(DATA / "sadeeddiac-25.parquet") + inputs = [DIACRITICS_RE.sub("", t) for t in + table.column("input").to_pylist()] + outputs = table.column("output").to_pylist() + windows, counts = [], [] + for text in inputs: + ws = split_windows(text) + counts.append(len(ws)) + windows.extend(ws) + preds = [] + with torch.no_grad(): + for i in range(0, len(windows), 32): + enc = tok(windows[i:i + 32], 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) + preds.extend(tok.batch_decode(gen, skip_special_tokens=True)) + k = 0 + paragraphs = [] + for text, c in zip(inputs, counts): + paragraphs.append(project_haraqat( + " ".join(preds[k:k + c]), text)) + k += c + csv_path = Path("/tmp/sadeed_oracle.csv") + pd.DataFrame({"gt": outputs, "pred": paragraphs}).to_csv( + csv_path, index=False, header=False) + from sadeed_evaluator import ArabicDiacritizationEvaluator as E + + print("===== SadeedDiac-25 (default protocol) =====", flush=True) + E.report_errors_on_csv_file(str(csv_path), ground_truth_column_index=0, + predicted_column_index=1, has_header=False, + gt_missing_diacritic_is_error=False) + + import eval_wikinews_multiref as WN + + bench = [l for l in (DATA / "wikinews/WikiNews_2024_Multi_Ref.txt.diac") + .read_text(encoding="utf-8").splitlines() 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): + enc = tok(wwindows[i:i + 32], 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)) + k = 0 + wparas = [] + for text, c in zip(winputs, wcounts): + wparas.append(project_haraqat(" ".join(wpreds[k:k + c]), text)) + k += c + wscore = WN.score(wparas, bench, skip_last=False) + print(json.dumps({"run": CKPT.name, "surface": + "wikinews-2024-multiref", **wscore}, + ensure_ascii=False), flush=True) + fa = [" ".join(w.split("/", 1)[0] for w in p.split()) for p in wparas] + fascore = WN.score(fa, bench, skip_last=False) + verdict = {"run": CKPT.name, "surface": "wikinews-firstalt", + "gate": "raw wer < 17.3794", "r7_shipped": 17.3794, + "stretch": 10.1262, **wscore} + print(json.dumps(verdict, ensure_ascii=False), flush=True) + (CKPT / "oracle_verdict.json").write_text( + json.dumps({"sadeed_csv": "written", "wikinews_raw": wscore, + "wikinews_firstalt": verdict}, ensure_ascii=False, + indent=1)) + (CKPT / "TRAIN_DONE").write_text("ok\n") + + +if __name__ == "__main__": + main() From b974b701286fa46a96329653c62e7b4b98b2e5cc Mon Sep 17 00:00:00 2001 From: Ronald Tse Date: Fri, 30 Oct 2026 18:09:00 +0800 Subject: [PATCH 2/5] feat(WO64): news-arm decode modes + tie-skip + rotating parts Stage-1 refinements from the smoke runs: rotating part files with row-keyed resume (mount only flushes closed files; the jobs-token bucket API 401s), tie-skip of rows where the primary arm exactly matches silver (news cannot beat a tie), and D_NEWS_MODE gpu|cpu|cpu_par with a spawn-based CPU pool (the int8 plain arm is fastest as parallel CPU single-row; forked children must not inherit the parent's CUDA context). Consolidation handles skipped rows and reports per-arm evaluated-row CER. --- oracle_relabel_onnx.py | 199 ++++++++++++++++++++++++++++++----------- 1 file changed, 145 insertions(+), 54 deletions(-) diff --git a/oracle_relabel_onnx.py b/oracle_relabel_onnx.py index 495f3ea..dc2fc2b 100644 --- a/oracle_relabel_onnx.py +++ b/oracle_relabel_onnx.py @@ -90,12 +90,15 @@ def lev(a: str, b: str) -> int: class Arm: - def __init__(self, name: str, zip_path: str): + def __init__(self, name: str, zip_path: str, force_cpu: bool = False, + intra_threads: int = 0): import onnxruntime as ort self.name = name opts = ort.SessionOptions() opts.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL + if intra_threads: + opts.intra_op_num_threads = intra_threads graphs, digests = {}, {} with zipfile.ZipFile(zip_path) as z: for member in z.namelist(): @@ -122,11 +125,15 @@ def __init__(self, name: str, zip_path: str): raise RuntimeError(f"{zip_path}:{fname} sha256 mismatch") providers = [p for p in ("CUDAExecutionProvider", "CPUExecutionProvider") if p in ort.get_available_providers()] + if force_cpu: + providers = ["CPUExecutionProvider"] self.enc = ort.InferenceSession(graphs["encoder.onnx"], opts, providers) dname = ("decoder-kv.onnx" if "decoder-kv.onnx" in graphs else "decoder.onnx") self.kv = "decoder-kv.onnx" in graphs self.dec = ort.InferenceSession(graphs[dname], opts, providers) + print(f"[{name}] EP={self.enc.get_providers()} kv={self.kv}", + flush=True) self.out_names = [o.name for o in self.dec.get_outputs()] self.pasts = {} for m in self.dec.get_inputs(): @@ -272,83 +279,136 @@ def decode_groups(nbytes: list[int], done: set[int], return groups -def decode_arm(arm: Arm, srcs: list[str]) -> None: - raw = CKPT / f"preds-{arm.name}.raw.jsonl" +def decode_arm(arm: Arm, srcs: list[str], skip: set[int] | None = None, + pool=None) -> None: + """Rotating part files (the khm-labeler pattern): the mount only + reliably flushes closed files, and HfFileSystem API sync 401s from + inside jobs — so durability = new part file every CHUNK rows. + + skip: rows that need no decode (news cannot beat an exact r7 match; + ties go to r7) — they are treated as decoded for planning/resume. + pool: optional multiprocessing pool of CPU workers (int8 plain arm + is fastest as parallel CPU single-row; see the bench verdict). + """ + CHUNK = 2000 + skip = skip or set() + parts_dir = CKPT / f"parts-{arm.name}" + parts_dir.mkdir(parents=True, exist_ok=True) done: set[int] = set() - if raw.exists(): - for line in raw.read_text(encoding="utf-8").splitlines(): + for part in sorted(parts_dir.glob("part-*.jsonl")): + for line in part.read_text(encoding="utf-8").splitlines(): if line.strip(): done.add(json.loads(line)["i"]) - print(f"[{arm.name}] resume: {len(done)} rows already decoded", + n_decodable = len(srcs) - len(skip) + if done: + print(f"[{arm.name}] resume: {len(done)}/{n_decodable} rows decoded", flush=True) nbytes = [len(s.encode()) + 1 for s in srcs] # plain-decoder exports are batch=1-traced (parity test: batched # plain decode corrupts; only the kv graph is batch-safe) — so the # int8 news arm always decodes single-row, exactly like the runtime - groups = decode_groups(nbytes, done, BS if arm.kv else 1) + groups = decode_groups(nbytes, done | skip, BS if arm.kv else 1) t0 = time.time() n_done = 0 - with raw.open("a", encoding="utf-8") as f: - for gi, g in enumerate(groups): - preds = arm.decode_equal([srcs[i] for i in g]) - for i, p in zip(g, preds): - f.write(json.dumps({"i": i, "p": p}, ensure_ascii=False) + "\n") - n_done += len(g) - f.flush() - if gi % 40 == 0: - try: - sync_to_bucket(raw) - except Exception as e: # noqa: BLE001 - print(f"[{arm.name}] sync failed: {e}", flush=True) - if gi % 20 == 0: + part_idx = len(list(parts_dir.glob("part-*.jsonl"))) + f = (parts_dir / f"part-{part_idx:03d}.jsonl").open( + "w", encoding="utf-8") + + def write(i: int, p: str) -> None: + nonlocal n_done, part_idx, f + f.write(json.dumps({"i": i, "p": p}, ensure_ascii=False) + "\n") + n_done += 1 + if n_done % CHUNK == 0: + f.close() + part_idx += 1 + f = (parts_dir / f"part-{part_idx:03d}.jsonl").open( + "w", encoding="utf-8") + print(f"[{arm.name}] part {part_idx - 1:03d} flushed", + flush=True) + + if pool is not None: + jobs = [(i, srcs[i]) for g in groups for i in g] + for i, p in pool.imap_unordered(_pool_worker, jobs, chunksize=4): + write(i, p) + if n_done % 100 == 0: r = n_done / max(1, time.time() - t0) - rem = len(srcs) - len(done) - n_done - print(f"[{arm.name}] {n_done}/{len(srcs) - len(done)} " - f"({r:.0f} rows/s, eta {rem / max(r, 1) / 3600:.1f}h)", - flush=True) - try: - sync_to_bucket(raw) - except Exception as e: # noqa: BLE001 - print(f"[{arm.name}] final sync failed: {e}", flush=True) - - -def raw_preds(name: str, n: int) -> list[str] | None: - raw = CKPT / f"preds-{name}.raw.jsonl" - if not raw.exists(): - return None + print(f"[{arm.name}] {n_done}/{n_decodable} " + f"({r:.1f} rows/s)", flush=True) + f.close() + print(f"[{arm.name}] done ({n_done} this session)", flush=True) + return + + for gi, g in enumerate(groups): + preds = arm.decode_equal([srcs[i] for i in g]) + for i, p in zip(g, preds): + write(i, p) + if gi % 20 == 0: + r = n_done / max(1, time.time() - t0) + rem = n_decodable - len(done) - n_done + print(f"[{arm.name}] {n_done}/{n_decodable} " + f"({r:.1f} rows/s, eta {rem / max(r, 1) / 3600:.1f}h)", + flush=True) + f.close() + print(f"[{arm.name}] done ({n_done} this session)", flush=True) + + +_WORKER_ARM = None + + +def _pool_worker(job: tuple[int, str]) -> tuple[int, str]: + global _WORKER_ARM + if _WORKER_ARM is None: + _WORKER_ARM = Arm("news", ARMS[-1][1], force_cpu=True, + intra_threads=2) + i, src = job + return i, _WORKER_ARM.decode_equal([src])[0] + + +def raw_preds(name: str, n: int, skip: set[int] | None = None): + skip = skip or set() + parts_dir = CKPT / f"parts-{name}" by_i: dict[int, str] = {} - for line in raw.read_text(encoding="utf-8").splitlines(): - if line.strip(): - r = json.loads(line) - by_i[r["i"]] = r["p"] - if len(by_i) < n: - print(f"[{name}] raw covers {len(by_i)}/{n} rows", flush=True) + for part in sorted(parts_dir.glob("part-*.jsonl")): + for line in part.read_text(encoding="utf-8").splitlines(): + if line.strip(): + r = json.loads(line) + by_i[r["i"]] = r["p"] + have = len(by_i) + len(skip - set(by_i)) + if have < n: + print(f"[{name}] raw covers {have}/{n} rows", flush=True) return None - return [by_i[i] for i in range(n)] + return [by_i.get(i) for i in range(n)] def consolidate(rows: list[tuple[str, str]], - preds: dict[str, list[str]]) -> None: + preds: dict[str, list[str | None]]) -> None: with (CKPT / "arm-preds.jsonl").open("w", encoding="utf-8") as f, \ (CKPT / "oracle-targets.jsonl").open("w", encoding="utf-8") as g: cers = {n: 0.0 for n, _ in ARMS} + eval_n = {n: 0 for n, _ in ARMS} oracle_cer = 0.0 disagree = 0 for i, (src, silver) in enumerate(rows): row_preds = {n: preds[n][i] for n, _ in ARMS} - if len(set(row_preds.values())) > 1: + if len({p for p in row_preds.values() if p is not None}) > 1: disagree += 1 - best_name, best_d = ARMS[0][0], lev(row_preds[ARMS[0][0]], silver) - best_tgt = row_preds[ARMS[0][0]] + primary = ARMS[0][0] + best_name, best_d = primary, lev(row_preds[primary], silver) + best_tgt = row_preds[primary] for n, _ in ARMS[1:]: + if row_preds[n] is None: # skipped: cannot beat the tie + continue d = lev(row_preds[n], silver) if d < best_d: best_name, best_d, best_tgt = n, d, row_preds[n] for n in cers: + if row_preds[n] is None: + continue cers[n] += lev(row_preds[n], silver) / max(1, len(silver)) + eval_n[n] += 1 oracle_cer += best_d / max(1, len(silver)) - f.write(json.dumps({"src": src, "silver": silver, **row_preds}, - ensure_ascii=False) + "\n") + f.write(json.dumps({"src": src, "silver": silver, + **row_preds}, ensure_ascii=False) + "\n") g.write(json.dumps({"src": src, "tgt": best_tgt, "pick": best_name}, ensure_ascii=False) + "\n") n = len(rows) @@ -358,7 +418,9 @@ def consolidate(rows: list[tuple[str, str]], p = json.loads(l)["pick"] picks[p] = picks.get(p, 0) + 1 stats = {"rows": n, "disagreement": round(disagree / n, 4), - "arm_cer": {k: round(v / n, 4) for k, v in cers.items()}, + "arm_cer": {k: round(v / max(1, eval_n[k]), 4) + for k, v in cers.items()}, + "arm_eval_rows": eval_n, "oracle_cer": round(oracle_cer / n, 4), "picks": picks} print(json.dumps(stats), flush=True) (CKPT / "relabel_stats.json").write_text(json.dumps(stats, indent=1)) @@ -368,12 +430,41 @@ def main() -> None: CKPT.mkdir(parents=True, exist_ok=True) rows = build_rows() srcs = [s for s, _ in rows] - loaded = [(n, Arm(n, p)) for n, p in ARMS] - for name, arm in loaded: - decode_arm(arm, srcs) - preds = {} - for name, _ in loaded: - p = raw_preds(name, len(rows)) + silvers = [t for _, t in rows] + primary_name, primary_zip = ARMS[0] + primary = Arm(primary_name, primary_zip) + decode_arm(primary, srcs) + p_primary = raw_preds(primary_name, len(rows)) + if p_primary is None: + print(f"[consolidate] waiting on arm {primary_name}", flush=True) + return + skip = {i for i in range(len(rows)) if lev(p_primary[i], silvers[i]) == 0} + print(f"[skip] {len(skip)}/{len(rows)} rows exact under {primary_name} " + f"— news cannot beat a tie; skipped", flush=True) + (CKPT / "skip.json").write_text(json.dumps(sorted(skip))) + mode = os.environ.get("D_NEWS_MODE", "gpu") + if mode == "cpu": + news = Arm(ARMS[-1][0], ARMS[-1][1], force_cpu=True, + intra_threads=4) + else: + news = Arm(ARMS[-1][0], ARMS[-1][1]) + pool = None + if mode == "cpu_par": + import multiprocessing as mp + + # spawn, not fork: the parent already initialized the CUDA + # context for the primary arm; forked children must not inherit it + pool = mp.get_context("spawn").Pool( + int(os.environ.get("D_NEWS_PROCS", "8"))) + try: + decode_arm(news, srcs, skip=skip, pool=pool) + finally: + if pool is not None: + pool.close() + pool.join() + preds = {primary_name: p_primary} + for name, _ in ARMS[1:]: + p = raw_preds(name, len(rows), skip=skip) if p is None: print(f"[consolidate] waiting on arm {name}", flush=True) return From ff8eb8bb484a369000c7697f059e6b9ce9fcb55c Mon Sep 17 00:00:00 2001 From: Ronald Tse Date: Fri, 30 Oct 2026 18:29:33 +0800 Subject: [PATCH 3/5] feat(WO64): env-driven second arm + two-part fetch from Releases MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit D_ARM2_NAME/D_ARM2_ZIP/D_ARM2_PARTS assemble the second oracle arm from sha256-pinned GitHub Release parts (canonical data channel). Launch arm set = r7 (ara-diac-2.0 kv) + r6 (ara-diac-1.0 kv, parity-verified vs the house runtime) — both batchable; the int8 news arm decodes at ~1 row/min GPU / 23 s/row CPU and is unusable for corpus-scale relabeling. --- oracle_relabel_onnx.py | 37 ++++++++++++++++++++++++++++++++++++- 1 file changed, 36 insertions(+), 1 deletion(-) diff --git a/oracle_relabel_onnx.py b/oracle_relabel_onnx.py index dc2fc2b..3cd9f09 100644 --- a/oracle_relabel_onnx.py +++ b/oracle_relabel_onnx.py @@ -49,7 +49,8 @@ BS = int(os.environ.get("D_BS", "48")) ARMS = [ ("r7", "/ckpt/b3-arms/ara-diac-2.0-fp32.zip"), - ("news", "/ckpt/b3-arms/ara-diac-news-1.0.zip"), + (os.environ.get("D_ARM2_NAME", "news"), + os.environ.get("D_ARM2_ZIP", "/ckpt/b3-arms/ara-diac-news-1.0.zip")), ] DIACRITICS_RE = re.compile("[ؐ-ًؚ-ٰٟۖ-ۜ۟-۪ۨ-ۭ]") @@ -428,6 +429,40 @@ def consolidate(rows: list[tuple[str, str]], def main() -> None: CKPT.mkdir(parents=True, exist_ok=True) + arm2_url = os.environ.get("D_ARM2_URL") + parts_spec = os.environ.get("D_ARM2_PARTS") # url|sha,url|sha,... + if parts_spec: + import hashlib + import urllib.request + + dst = Path(ARMS[-1][1]) + if not dst.exists(): + with dst.open("wb") as out: + for spec in parts_spec.split(","): + url, want = spec.rsplit("|", 1) + print(f"[fetch] {url}", flush=True) + with urllib.request.urlopen(url) as r: + blob = r.read() + got = hashlib.sha256(blob).hexdigest() + if got != want: + raise RuntimeError( + f"part sha256 mismatch {got} != {want}") + out.write(blob) + print(f"[fetch] arm2 assembled {dst} ({dst.stat().st_size} B)", + flush=True) + elif arm2_url: + import hashlib + import urllib.request + + dst = Path(ARMS[-1][1]) + want = os.environ.get("D_ARM2_SHA", "") + if not dst.exists(): + print(f"[fetch] {arm2_url}", flush=True) + urllib.request.urlretrieve(arm2_url, dst) + got = hashlib.sha256(dst.read_bytes()).hexdigest() + if want and got != want: + raise RuntimeError(f"arm2 sha256 mismatch {got} != {want}") + print(f"[fetch] arm2 sha256 OK ({got[:12]})", flush=True) rows = build_rows() srcs = [s for s, _ in rows] silvers = [t for _, t in rows] From 1e60054a23ff51d335c649c3ca24f5a3dadec8c1 Mon Sep 17 00:00:00 2001 From: Ronald Tse Date: Sat, 31 Oct 2026 02:12:25 +0800 Subject: [PATCH 4/5] feat(WO64): per-arm parallel jobs (D_ARM_ONLY) + decode-length cap (D_MAXB) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The full-550B corpus needs 1.5-2 GPU-days per arm at measured rates (length-sorted ascending means the tail is quadratic). D_MAXB=300 caps the decode length (~85% of rows, a fraction of the cost, and every row fits stage-2's 600B window with no stitching); D_ARM_ONLY runs each arm as its own job — parts-/ are independent, consolidation runs in whichever job sees both complete. Resume parts are plan-indexed: a changed corpus recipe (D_SUBSET/D_MAXB) MUST relaunch under a fresh D_RUN — mixing indices across plans corrupts the oracle. --- oracle_relabel_onnx.py | 89 ++++++++++++++++++++++++------------------ 1 file changed, 50 insertions(+), 39 deletions(-) diff --git a/oracle_relabel_onnx.py b/oracle_relabel_onnx.py index 3cd9f09..4fed8a4 100644 --- a/oracle_relabel_onnx.py +++ b/oracle_relabel_onnx.py @@ -218,7 +218,11 @@ def build_rows() -> list[tuple[str, str]]: 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: + # D_MAXB caps the decode length (default 550 = the recorded + # recipe; 300 was used for tractable corpus-scale relabeling — + # recorded in relabel_stats) + if (len(src) >= 5 and len(src.encode()) <= WINDOW - 50 + and len(src.encode()) <= int(os.environ.get("D_MAXB", "550"))): pseudo.append((src, tgt)) for aline in articles: @@ -419,6 +423,7 @@ def consolidate(rows: list[tuple[str, str]], p = json.loads(l)["pick"] picks[p] = picks.get(p, 0) + 1 stats = {"rows": n, "disagreement": round(disagree / n, 4), + "maxb": int(os.environ.get("D_MAXB", "550")), "arm_cer": {k: round(v / max(1, eval_n[k]), 4) for k, v in cers.items()}, "arm_eval_rows": eval_n, @@ -466,45 +471,51 @@ def main() -> None: rows = build_rows() srcs = [s for s, _ in rows] silvers = [t for _, t in rows] - primary_name, primary_zip = ARMS[0] - primary = Arm(primary_name, primary_zip) - decode_arm(primary, srcs) - p_primary = raw_preds(primary_name, len(rows)) - if p_primary is None: - print(f"[consolidate] waiting on arm {primary_name}", flush=True) - return - skip = {i for i in range(len(rows)) if lev(p_primary[i], silvers[i]) == 0} - print(f"[skip] {len(skip)}/{len(rows)} rows exact under {primary_name} " - f"— news cannot beat a tie; skipped", flush=True) - (CKPT / "skip.json").write_text(json.dumps(sorted(skip))) - mode = os.environ.get("D_NEWS_MODE", "gpu") - if mode == "cpu": - news = Arm(ARMS[-1][0], ARMS[-1][1], force_cpu=True, - intra_threads=4) - else: - news = Arm(ARMS[-1][0], ARMS[-1][1]) - pool = None - if mode == "cpu_par": - import multiprocessing as mp - - # spawn, not fork: the parent already initialized the CUDA - # context for the primary arm; forked children must not inherit it - pool = mp.get_context("spawn").Pool( - int(os.environ.get("D_NEWS_PROCS", "8"))) - try: - decode_arm(news, srcs, skip=skip, pool=pool) - finally: - if pool is not None: - pool.close() - pool.join() - preds = {primary_name: p_primary} - for name, _ in ARMS[1:]: - p = raw_preds(name, len(rows), skip=skip) - if p is None: - print(f"[consolidate] waiting on arm {name}", flush=True) + arm_only = os.environ.get("D_ARM_ONLY") + skip: set[int] = set() + + if not arm_only or arm_only == ARMS[0][0]: + primary = Arm(ARMS[0][0], ARMS[0][1]) + decode_arm(primary, srcs) + + if not arm_only: + # sequential mode: the tie-skip needs the primary's full preds + p0 = raw_preds(ARMS[0][0], len(rows)) + if p0 is None: + print(f"[consolidate] waiting on arm {ARMS[0][0]}", flush=True) return - preds[name] = p - consolidate(rows, preds) + skip = {i for i in range(len(rows)) + if lev(p0[i], silvers[i]) == 0} + print(f"[skip] {len(skip)}/{len(rows)} rows exact under " + f"{ARMS[0][0]} — arm2 cannot beat a tie; skipped", flush=True) + (CKPT / "skip.json").write_text(json.dumps(sorted(skip))) + + if not arm_only or arm_only != ARMS[0][0]: + second = Arm(ARMS[-1][0], ARMS[-1][1]) + pool = None + if os.environ.get("D_NEWS_MODE") == "cpu_par": + import multiprocessing as mp + + # spawn, not fork: the parent may hold a CUDA context + pool = mp.get_context("spawn").Pool( + int(os.environ.get("D_NEWS_PROCS", "8"))) + try: + decode_arm(second, srcs, skip=skip, pool=pool) + finally: + if pool is not None: + pool.close() + pool.join() + + # consolidation runs in whichever job sees both arms complete + p0 = raw_preds(ARMS[0][0], len(rows)) + p1 = raw_preds(ARMS[-1][0], len(rows), skip=skip) + if p0 is None: + print(f"[consolidate] waiting on arm {ARMS[0][0]}", flush=True) + return + if p1 is None: + print(f"[consolidate] waiting on arm {ARMS[-1][0]}", flush=True) + return + consolidate(rows, {ARMS[0][0]: p0, ARMS[-1][0]: p1}) (CKPT / "RELABEL_DONE").write_text("ok\n") try: sync_to_bucket(CKPT / "RELABEL_DONE") From 9f5b04e76681d2e62e8b1d4f55d747d3c6f7b0d3 Mon Sep 17 00:00:00 2001 From: Ronald Tse Date: Sat, 31 Oct 2026 21:59:29 +0800 Subject: [PATCH 5/5] feat(WO65): env-driven primary arm + D_ROWS_JSONL pre-paired corpus mode Generalizes the oracle relabel beyond the Arabic qcri recipe: D_ARM1_NAME/ D_ARM1_ZIP select the primary arm from env, and D_ROWS_JSONL reads a pre-paired {src,tgt} jsonl (stripped, seed-42 shuffled, subset) instead of building rows from article text. This makes the Urdu WO65 campaign (arms urd-diac-1.0 + urd-diac-base-1.0 over the 573K-row urd-dia train split, gate CER < 3) a single launch command once GPU budget frees. --- oracle_relabel_onnx.py | 30 ++++++++++++++++++++++++++++-- 1 file changed, 28 insertions(+), 2 deletions(-) diff --git a/oracle_relabel_onnx.py b/oracle_relabel_onnx.py index 4fed8a4..2109e58 100644 --- a/oracle_relabel_onnx.py +++ b/oracle_relabel_onnx.py @@ -48,7 +48,8 @@ SMOKE = bool(os.environ.get("D_SMOKE")) BS = int(os.environ.get("D_BS", "48")) ARMS = [ - ("r7", "/ckpt/b3-arms/ara-diac-2.0-fp32.zip"), + (os.environ.get("D_ARM1_NAME", "r7"), + os.environ.get("D_ARM1_ZIP", "/ckpt/b3-arms/ara-diac-2.0-fp32.zip")), (os.environ.get("D_ARM2_NAME", "news"), os.environ.get("D_ARM2_ZIP", "/ckpt/b3-arms/ara-diac-news-1.0.zip")), ] @@ -208,7 +209,32 @@ def split_windows(text: str, budget: int = WINDOW) -> list[str]: def build_rows() -> list[tuple[str, str]]: - """run-047's exact qcri-only silver corpus, stripped, shuffled.""" + """run-047's exact qcri-only silver corpus, stripped, shuffled. + + D_ROWS_JSONL switches to a pre-paired corpus: a jsonl of {"src", + "tgt"} rows (e.g. the urd-dia train split), stripped of diacritics + on the src side, seed-42 shuffled, subset taken after the shuffle. + """ + rows_path = os.environ.get("D_ROWS_JSONL") + if rows_path: + pairs: list[tuple[str, str]] = [] + maxb = int(os.environ.get("D_MAXB", "550")) + for line in Path(rows_path).read_text(encoding="utf-8").splitlines(): + if not line.strip(): + continue + r = json.loads(line) + tgt = str(r.get("tgt", "")).strip() + src = DIACRITICS_RE.sub("", str(r.get("src", ""))).strip() + if len(src) >= 5 and len(src.encode()) <= maxb: + pairs.append((src, tgt)) + if len(pairs) >= PSEUDO_CAP: + break + rng = random.Random(42) + rng.shuffle(pairs) + rows = pairs[:SUBSET] + print(f"[data] jsonl rows={len(pairs)} subset={len(rows)}", + flush=True) + return rows qpath = DATA / "qcri-wiki.jsonl" articles = qpath.read_text(encoding="utf-8").splitlines() rng2 = random.Random(7)