diff --git a/oracle_relabel_onnx.py b/oracle_relabel_onnx.py new file mode 100644 index 0000000..2109e58 --- /dev/null +++ b/oracle_relabel_onnx.py @@ -0,0 +1,562 @@ +"""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 = [ + (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")), +] +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, 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(): + 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()] + 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(): + 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. + + 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) + rng2.shuffle(articles) + pseudo: list[tuple[str, str]] = [] + + def add_silver(tgt: str) -> None: + tgt = tgt.strip() + src = DIACRITICS_RE.sub("", tgt).strip() + # 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: + 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], 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() + 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"]) + 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 | skip, BS if arm.kv else 1) + t0 = time.time() + n_done = 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) + 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 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.get(i) for i in range(n)] + + +def consolidate(rows: list[tuple[str, str]], + 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({p for p in row_preds.values() if p is not None}) > 1: + disagree += 1 + 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") + 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), + "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, + "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) + 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] + 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 + 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") + 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()