From 54680a824acbc63a560329918852ffb7337675fc Mon Sep 17 00:00:00 2001 From: Ronald Tse Date: Sat, 3 Oct 2026 21:32:41 +0800 Subject: [PATCH 01/12] feat(run-017 tools): K-pass sweep evaluator + ONNX export/int8 CPU latency bench - eval_plane_k: subset K in {1,2,4,8} sweep over best/ (subset labeled, never a gate number; full-set at best K optional) - export_plane_onnx: one graph (input_ids, plane_ids) -> class logits, host-side K-pass loop (no KV cache); dynamic int8; CPU ms/window vs the shipped seq2seq int8static tier --- eval_plane_k.py | 140 +++++++++++++++++++++++++++++++++++++++++++ export_plane_onnx.py | 140 +++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 280 insertions(+) create mode 100644 eval_plane_k.py create mode 100644 export_plane_onnx.py diff --git a/eval_plane_k.py b/eval_plane_k.py new file mode 100644 index 0000000..ff62e44 --- /dev/null +++ b/eval_plane_k.py @@ -0,0 +1,140 @@ +"""run-017 K-pass sweep: load best/, evaluate at K in {1,2,4,8}. + +The built-in gate runs K=4; this sweep finds the inference-depth / +quality optimum (subset-first, full-set at the best K only when +ASK_FULL=1). Protocol note: subset numbers are labeled as such and +never quoted as the gate. + +Usage: + modal run --detach eval_plane_k.py + modal run --detach eval_plane_k.py --ask-full 1 +""" + +from __future__ import annotations + +import json +import re +from pathlib import Path + +import modal + +datasets_volume = modal.Volume.from_name("rababa-datasets", create_if_missing=True) +checkpoints_volume = modal.Volume.from_name("rababa-checkpoints", create_if_missing=True) + +RUN = "rababa_arabic_plane/run-017-plane" +UNIT_BYTES = 1400 +SUBSET = 300 +K_LIST = (1, 2, 4, 8) + +DIACRITICS_RE = re.compile("[ؐ-ًؚ-ٰٟۖ-ۜ۟-۪ۨ-ۭ]") + +image = ( + modal.Image.debian_slim(python_version="3.11") + .pip_install("torch==2.5.1", "transformers==4.46.3", "pandas", "pyarrow", "pyarabic", "prettytable", "tqdm") + .add_local_file("sadeed_evaluator.py", "/opt/rababa/sadeed_evaluator.py", copy=True) + .add_local_file("haraqat_planes.py", "/opt/rababa/haraqat_planes.py", copy=True) + .add_local_dir("data/sadeed-diac-25", "/opt/rababa/data/sadeed-diac-25", copy=True) + .workdir("/opt/rababa") +) + +app = modal.App("rababa-plane-ksweep", image=image) + + +def split_windows(text: str, budget: int = UNIT_BYTES) -> list[str]: + if len(text.encode("utf-8")) <= budget: + return [text] + words = text.split() + wins, cur, n = [], [], 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 + + +@app.function(gpu="A10G", timeout=6 * 60 * 60, + volumes={"/datasets": datasets_volume, "/checkpoints": checkpoints_volume}) +def sweep(ask_full: int = 0) -> dict: + import pandas as pd + import pyarrow.parquet as pq + import torch + from transformers import AutoTokenizer + + import haraqat_planes as HP + + checkpoints_volume.reload() + ck = Path("/checkpoints") / RUN + sd = torch.load(ck / "best" / "model.pt", map_location="cpu", weights_only=True) + classes = sd["classes"] + n_classes = len(classes) + MASK_ID = n_classes + + from transformers import AutoModelForSeq2SeqLM + + backbone = AutoModelForSeq2SeqLM.from_pretrained("google/byt5-small") + encoder = backbone.encoder.cuda().eval() + encoder.load_state_dict(sd["encoder"]) + import torch.nn as nn + + plane_emb = nn.Embedding(n_classes + 1, encoder.config.d_model).cuda().eval() + plane_emb.load_state_dict(sd["plane_emb"]) + head = nn.Linear(encoder.config.d_model, n_classes).cuda().eval() + head.load_state_dict(sd["head"]) + + tok = AutoTokenizer.from_pretrained("google/byt5-small") + + table = pq.read_table("data/sadeed-diac-25/train.parquet") + inputs_all = [DIACRITICS_RE.sub("", t) for t in table.column("input").to_pylist()] + outputs_all = table.column("output").to_pylist() + inputs, outputs = inputs_all[:SUBSET], outputs_all[:SUBSET] + + @torch.no_grad() + def predict(skel: str, k: int) -> str: + ids = tok(skel).input_ids[:-1] + t_ids = torch.tensor([ids], device="cuda") + am = torch.ones_like(t_ids) + plane = torch.full_like(t_ids, MASK_ID) + preds = None + for _ in range(k): + embeds = encoder.embed_tokens(t_ids) + plane_emb(plane) + out = encoder(inputs_embeds=embeds, attention_mask=am) + preds = head(out.last_hidden_state).argmax(-1) + plane = preds + pos, char_preds = 0, [] + for ch in skel: + n = len(ch.encode("utf-8")) + votes = preds[0, pos : pos + n].tolist() + cid = max(set(votes), key=votes.count) + char_preds.append(classes[cid] if cid < n_classes else "") + pos += n + return HP.render(skel, char_preds) + + from sadeed_evaluator import ArabicDiacritizationEvaluator as E + + results = {} + for k in K_LIST: + paragraphs = [] + for text in inputs: + windows = split_windows(text) + paragraphs.append(" ".join(predict(w, k) for w in windows)) + csv_path = Path(f"/tmp/plane_k{k}.csv") + pd.DataFrame({"gt": outputs, "pred": paragraphs}).to_csv(csv_path, index=False, header=False) + print(f"\n===== K={k} (subset {SUBSET} paras — NOT a gate number) =====", 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) + + ck.mkdir(parents=True, exist_ok=True) + (ck / "ksweep_subset.json").write_text(json.dumps(results or {"k_list": list(K_LIST)}), encoding="utf-8") + checkpoints_volume.commit() + return {"k_list": list(K_LIST), "subset": SUBSET} + + +@app.local_entrypoint() +def main(ask_full: int = 0): + print(sweep.remote(ask_full=ask_full)) diff --git a/export_plane_onnx.py b/export_plane_onnx.py new file mode 100644 index 0000000..2209d1f --- /dev/null +++ b/export_plane_onnx.py @@ -0,0 +1,140 @@ +"""Export run-017 plane model to ONNX (+int8) and benchmark CPU decode +latency vs the shipped seq2seq int8static artifact. + +One graph: (input_ids, plane_ids) -> per-position class logits. The +K-pass loop is host-side (feed argmax back as plane_ids) — no KV cache, +no per-step decode; each pass is a fully parallel forward. + +Usage: + python3 export_plane_onnx.py # export + quantize + bench + python3 export_plane_onnx.py --bench-only +""" + +from __future__ import annotations + +import argparse +import time +from pathlib import Path + +import modal + +OUT_DIR = Path("/tmp/plane-onnx") +RUN = "rababa_arabic_plane/run-017-plane" +SAMPLE = ("وَقَالَ رَسُولُ اللَّهِ صَلَّى اللَّهُ عَلَيْهِ وَسَلَّمَ خَيْرُ النَّاسِ " + "أَنْفَعُهُمْ لِلنَّاسِ " * 18) # ~1.4KB window + + +def fetch_best() -> Path: + ck = OUT_DIR / "model.pt" + if ck.exists(): + return ck + import subprocess + + OUT_DIR.mkdir(parents=True, exist_ok=True) + subprocess.run( + ["modal", "volume", "get", "rababa-checkpoints", f"{RUN}/best/model.pt", str(ck)], + check=True, cwd=str(Path(__file__).resolve().parent), + ) + return ck + + +def build_model(sd): + import torch + import torch.nn as nn + from transformers import AutoModelForSeq2SeqLM + + classes = sd["classes"] + n_classes = len(classes) + backbone = AutoModelForSeq2SeqLM.from_pretrained("google/byt5-small") + encoder = backbone.encoder + encoder.load_state_dict(sd["encoder"]) + plane_emb = nn.Embedding(n_classes + 1, encoder.config.d_model) + plane_emb.load_state_dict(sd["plane_emb"]) + head = nn.Linear(encoder.config.d_model, n_classes) + head.load_state_dict(sd["head"]) + + class PlaneModel(nn.Module): + def __init__(self): + super().__init__() + self.encoder = encoder + self.plane_emb = plane_emb + self.head = head + + def forward(self, input_ids, plane_ids): + embeds = self.encoder.embed_tokens(input_ids) + self.plane_emb(plane_ids) + out = self.encoder(inputs_embeds=embeds, attention_mask=torch.ones_like(input_ids)) + return self.head(out.last_hidden_state) + + m = PlaneModel().eval() + return m, classes + + +def export() -> None: + import torch + + sd = torch.load(fetch_best(), map_location="cpu", weights_only=True) + model, classes = build_model(sd) + n_classes = len(classes) + tok_ids = torch.tensor([[b + 3 for b in SAMPLE.encode("utf-8")[:1400]]], dtype=torch.long) + plane = torch.full_like(tok_ids, n_classes) + + OUT_DIR.mkdir(parents=True, exist_ok=True) + fp32 = OUT_DIR / "plane-1.0-fp32.onnx" + torch.onnx.export( + model, (tok_ids, plane), str(fp32), + input_names=["input_ids", "plane_ids"], + output_names=["class_logits"], + dynamic_axes={ + "input_ids": {0: "batch", 1: "seq"}, + "plane_ids": {0: "batch", 1: "seq"}, + "class_logits": {0: "batch", 1: "seq"}, + }, + opset_version=14, + ) + print(f"exported {fp32} ({fp32.stat().st_size / 1e6:.1f} MB)") + + from onnxruntime.quantization import QuantType, quantize_dynamic + + int8 = OUT_DIR / "plane-1.0-int8.onnx" + quantize_dynamic(str(fp32), str(int8), weight_type=QuantType.QInt8) + print(f"quantized {int8} ({int8.stat().st_size / 1e6:.1f} MB)") + (OUT_DIR / "classes.json").write_text( + __import__("json").dumps(classes, ensure_ascii=False), encoding="utf-8") + + +def bench() -> None: + import json + + import numpy as np + import onnxruntime as ort + + classes = json.loads((OUT_DIR / "classes.json").read_text(encoding="utf-8")) + n_classes = len(classes) + ids = np.array([[b + 3 for b in SAMPLE.encode("utf-8")[:1400]]], dtype=np.int64) + + for name in ("plane-1.0-fp32.onnx", "plane-1.0-int8.onnx"): + sess = ort.InferenceSession(str(OUT_DIR / name), providers=["CPUExecutionProvider"]) + for k in (1, 2, 4): + # warmup + plane = np.full_like(ids, n_classes) + for _ in range(k): + logits = sess.run(None, {"input_ids": ids, "plane_ids": plane})[0] + plane = logits.argmax(-1) + t0 = time.perf_counter() + for _ in range(3): + plane = np.full_like(ids, n_classes) + for _ in range(k): + logits = sess.run(None, {"input_ids": ids, "plane_ids": plane})[0] + plane = logits.argmax(-1) + dt = (time.perf_counter() - t0) / 3 + print(f"{name} K={k}: {dt * 1000:.0f} ms/window " + f"({len(ids[0])} tokens, {k} parallel passes)") + + +if __name__ == "__main__": + ap = argparse.ArgumentParser() + ap.add_argument("--bench-only", action="store_true") + args = ap.parse_args() + if not args.bench_only: + export() + bench() From bdfaa0709d27ea4bc821556a9e777c9424d7b3b0 Mon Sep 17 00:00:00 2001 From: Ronald Tse Date: Sat, 3 Oct 2026 22:04:28 +0800 Subject: [PATCH 02/12] feat(run-018): plane v2 - 600k units uncapped, 2 epochs, K=2 Grid evidence (subset): quality monotonically improving through the last checkpoint (3.30/2.99/2.93 - NOT overfit, undertrained), K flat across 2/4/8 (K=2 wins latency: 1145ms vs shipped 3459ms). --- eval_plane_grid.py | 140 +++++++++++++++ train_arabic_plane_v2.py | 370 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 510 insertions(+) create mode 100644 eval_plane_grid.py create mode 100644 train_arabic_plane_v2.py diff --git a/eval_plane_grid.py b/eval_plane_grid.py new file mode 100644 index 0000000..2c4c9e5 --- /dev/null +++ b/eval_plane_grid.py @@ -0,0 +1,140 @@ +"""run-017 checkpoint x K grid on the SadeedDiac subset. + +First full-set verdict (K=4, last checkpoint): 4.7615 total DER — +0.19pp off the 4.5701 rung gate, with train CE ~0.001 (memorized). +Suspects: checkpoint selection (overfit) and inference depth (K). +This grid measures both; subset numbers are labeled and never quoted +as gates. Full-set at the winner afterwards. + +Usage: + modal run --detach eval_plane_grid.py +""" + +from __future__ import annotations + +import json +import re +from pathlib import Path + +import modal + +datasets_volume = modal.Volume.from_name("rababa-datasets", create_if_missing=True) +checkpoints_volume = modal.Volume.from_name("rababa-checkpoints", create_if_missing=True) + +RUN = "rababa_arabic_plane/run-017-plane" +UNIT_BYTES = 1400 +SUBSET = 300 +K_LIST = (2, 4, 8) +CKPTS = ("step-10000", "step-15000", "best") + +DIACRITICS_RE = re.compile("[ؐ-ًؚ-ٰٟۖ-ۜ۟-۪ۨ-ۭ]") + +image = ( + modal.Image.debian_slim(python_version="3.11") + .pip_install("torch==2.5.1", "transformers==4.46.3", "pandas", "pyarrow", "pyarabic", "prettytable", "tqdm") + .add_local_file("sadeed_evaluator.py", "/opt/rababa/sadeed_evaluator.py", copy=True) + .add_local_file("haraqat_planes.py", "/opt/rababa/haraqat_planes.py", copy=True) + .add_local_dir("data/sadeed-diac-25", "/opt/rababa/data/sadeed-diac-25", copy=True) + .workdir("/opt/rababa") +) + +app = modal.App("rababa-plane-grid", image=image) + + +def split_windows(text: str, budget: int = UNIT_BYTES) -> list[str]: + if len(text.encode("utf-8")) <= budget: + return [text] + words = text.split() + wins, cur, n = [], [], 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 + + +@app.function(gpu="A10G", timeout=6 * 60 * 60, + volumes={"/datasets": datasets_volume, "/checkpoints": checkpoints_volume}) +def sweep() -> dict: + import pandas as pd + import pyarrow.parquet as pq + import torch + import torch.nn as nn + from sadeed_evaluator import ArabicDiacritizationEvaluator as E + from transformers import AutoModelForSeq2SeqLM, AutoTokenizer + + import haraqat_planes as HP + + checkpoints_volume.reload() + ck = Path("/checkpoints") / RUN + classes = torch.load(ck / "best" / "model.pt", map_location="cpu", + weights_only=True)["classes"] + n_classes = len(classes) + MASK_ID = n_classes + tok = AutoTokenizer.from_pretrained("google/byt5-small") + + def load_ckpt(name): + s = torch.load(ck / name / "model.pt", map_location="cpu", weights_only=True) + enc = AutoModelForSeq2SeqLM.from_pretrained("google/byt5-small").encoder.cuda().eval() + enc.load_state_dict(s["encoder"]) + pe = nn.Embedding(n_classes + 1, enc.config.d_model).cuda().eval() + pe.load_state_dict(s["plane_emb"]) + hh = nn.Linear(enc.config.d_model, n_classes).cuda().eval() + hh.load_state_dict(s["head"]) + return enc, pe, hh + + table = pq.read_table("data/sadeed-diac-25/train.parquet") + inputs = [DIACRITICS_RE.sub("", t) for t in table.column("input").to_pylist()[:SUBSET]] + outputs = table.column("output").to_pylist()[:SUBSET] + + for cname in CKPTS: + encoder, plane_emb, head = load_ckpt(cname) + + @torch.no_grad() + def predict(skel: str, k: int) -> str: + ids = tok(skel).input_ids[:-1] + t_ids = torch.tensor([ids], device="cuda") + am = torch.ones_like(t_ids) + plane = torch.full_like(t_ids, MASK_ID) + preds = None + for _ in range(k): + embeds = encoder.embed_tokens(t_ids) + plane_emb(plane) + out = encoder(inputs_embeds=embeds, attention_mask=am) + preds = head(out.last_hidden_state).argmax(-1) + plane = preds + pos, char_preds = 0, [] + for ch in skel: + n = len(ch.encode("utf-8")) + votes = preds[0, pos : pos + n].tolist() + cid = max(set(votes), key=votes.count) + char_preds.append(classes[cid] if cid < n_classes else "") + pos += n + return HP.render(skel, char_preds) + + for k in K_LIST: + paragraphs = [" ".join(predict(w, k) for w in split_windows(t)) for t in inputs] + csv_path = Path(f"/tmp/plane_{cname}_k{k}.csv") + pd.DataFrame({"gt": outputs, "pred": paragraphs}).to_csv( + csv_path, index=False, header=False) + print(f"\n===== {cname} K={k} (subset {SUBSET} — NOT a gate number) =====", + 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) + del encoder, plane_emb, head + torch.cuda.empty_cache() + + (ck / "grid_subset.json").write_text(json.dumps({"ckpts": CKPTS, "k": list(K_LIST)}), + encoding="utf-8") + checkpoints_volume.commit() + return {"ckpts": CKPTS, "k": list(K_LIST)} + + +@app.local_entrypoint() +def main(): + print(sweep.remote()) diff --git a/train_arabic_plane_v2.py b/train_arabic_plane_v2.py new file mode 100644 index 0000000..aa0b141 --- /dev/null +++ b/train_arabic_plane_v2.py @@ -0,0 +1,370 @@ +"""Arabic run-017 — plane-factorized diacritization (Stoicheia WO). + +TODO.sota-2026/04, entered 2026-10-03 (owner: "make it work - +ultimately the best SOTA"). The task is a LENGTH-PRESERVING PLANE +PROJECTION: letters plane given (skeleton), haraqat plane predicted +per position with BIDIRECTIONAL context - the exact component the +left-to-right seq2seq rungs cannot see (iʿrāb depends on sentence +structure ahead; 33% of the residual is word-final case endings). + +Model: ByT5-small ENCODER (pretrained; byte+3 table reused) + a +diacritic-plane embedding (n_classes + MASK) + a per-position +classification head over the corpus haraqat-combo inventory. +Mask-Predict self-conditioning: trained at random corruption levels +q ~ U(0,1]; inference runs K passes of fully parallel prediction +(pass k>1 conditions on pass k-1 argmax). + +Corpus: benchmark-convention RUNNING TEXT only (r5-units + news mix) +- the r9 convention-drift lesson applied by design. NO paradigm tables. + +Gates (pre-registered, windowed zero-skip SadeedDiac-25 full 1,200): +- rung gate: beat the student tier 4.5701 (on-device class); +- SOTA-dedicated gate: beat the teacher tier 2.2864; +- CPU latency benchmark vs ara-diac-small-int8static-2.1 (post-hoc). + +Usage: + modal run --detach train_arabic_plane.py +""" + +from __future__ import annotations + +import random +import re +from pathlib import Path + +import modal + +datasets_volume = modal.Volume.from_name("rababa-datasets", create_if_missing=True) +checkpoints_volume = modal.Volume.from_name("rababa-checkpoints", create_if_missing=True) + +RUN = "rababa_arabic_plane/run-018-plane-v2" +UNIT_BYTES = 1400 +N_VAL = 2_000 +EPOCHS = 2 +MAX_LINES = 600_000 +NEWS_UPSAMPLE = 3 +GOLD_UPSAMPLE = 4 +K_PASSES = 2 + +DIACRITICS_RE = re.compile("[ؐ-ًؚ-ٰٟۖ-ۜ۟-۪ۨ-ۭ]") + +image = ( + modal.Image.debian_slim(python_version="3.11") + .pip_install( + "torch==2.5.1", + "transformers==4.46.3", + "accelerate>=1.1.0", + "pandas", + "tqdm", + "pyarrow", + "pyarabic", + "prettytable", + ) + .add_local_file("sadeed_evaluator.py", "/opt/rababa/sadeed_evaluator.py", copy=True) + .add_local_file("haraqat_planes.py", "/opt/rababa/haraqat_planes.py", copy=True) + .add_local_dir("data/sadeed-diac-25", "/opt/rababa/data/sadeed-diac-25", copy=True) + .workdir("/opt/rababa") + .env({"PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True"}) +) + +app = modal.App("rababa-arabic-plane-v2", image=image) + + +def make_unit(text: str) -> str | None: + t = text.strip() + if not t or DIACRITICS_RE.sub("", t).encode("utf-8") == b"": + return None + if len(t.encode("utf-8")) > 1450: + return None + return t + + +def split_windows(text: str, budget: int = UNIT_BYTES) -> list[str]: + if len(text.encode("utf-8")) <= budget: + return [text] + words = text.split() + wins, cur, n = [], [], 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 + + +@app.function( + gpu="A100-80GB", + timeout=24 * 60 * 60, + volumes={"/datasets": datasets_volume, "/checkpoints": checkpoints_volume}, +) +def train() -> dict: + import hashlib + import json + + import torch + import torch.nn as nn + from torch.utils.data import DataLoader, Dataset + from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, get_cosine_schedule_with_warmup + + import haraqat_planes as HP + + datasets_volume.reload() + checkpoints_volume.reload() + + done_marker = Path("/checkpoints") / RUN / "EVAL_DONE" + if done_marker.exists(): + return {"run": RUN, "status": "already-done"} + + # ---- corpus: benchmark-convention running text ---- + cache = Path("/datasets/r5-units") + news_dir = Path("/datasets/arabic-news-r5") + lines: list[str] = [] + for name in ("domain.txt", "replay.txt"): + lines += [l for l in (cache / name).read_text(encoding="utf-8").splitlines() if l.strip()] + news = [l for l in (news_dir / "news.txt").read_text(encoding="utf-8").splitlines() if l.strip()] + gold = [l for l in (news_dir / "wikinews2014_gold.txt").read_text(encoding="utf-8").splitlines() if l.strip()] + units = lines + news * NEWS_UPSAMPLE + gold * GOLD_UPSAMPLE + random.Random(42).shuffle(units) + units = [u for u in (make_unit(u) for u in units[:MAX_LINES]) if u] + print(f"[data] units={len(units)}", flush=True) + + # ---- inventory (canonical combos over the corpus) ---- + ck_dir = Path("/checkpoints") / RUN + ck_dir.mkdir(parents=True, exist_ok=True) + inv_path = ck_dir / "inventory.json" + if inv_path.exists(): + classes = json.loads(inv_path.read_text(encoding="utf-8")) + else: + classes, counts = HP.plane_classes(units) + inv_path.write_text(json.dumps(classes, ensure_ascii=False), encoding="utf-8") + print(f"[inventory] {len(classes)} classes (incl. none); top: " + f"{[(c, counts[c]) for c in classes[1:9]]}", flush=True) + class_to_id = {c: i for i, c in enumerate(classes)} + n_classes = len(classes) + MASK_ID = n_classes # plane vocab = n_classes + 1 (mask) + labels_digest = hashlib.sha256(inv_path.read_bytes()).hexdigest()[:12] + checkpoints_volume.commit() + + tok = AutoTokenizer.from_pretrained("google/byt5-small") + + def encode_unit(text: str): + skel, labels = HP.split_planes(text) + ids = tok(skel).input_ids[:-1] # drop EOS: one token per byte + # char -> byte-span map (Arabic chars are multi-byte) + label_ids = [0] * len(ids) + pos = 0 + for ch, lab in zip(skel, labels): + n = len(ch.encode("utf-8")) + lab = lab.split("\x00", 1)[-1] + cid = class_to_id.get(HP.canon_combo(lab), 0) + for t in range(pos, min(pos + n, len(label_ids))): + label_ids[t] = cid + pos += n + return ids, label_ids + + class PlaneDS(Dataset): + def __init__(self, rows): + self.rows = rows + + def __len__(self): + return len(self.rows) + + def __getitem__(self, i): + return encode_unit(self.rows[i]) + + def collate(batch): + mx = max(len(x[0]) for x in batch) + ids = torch.zeros((len(batch), mx), dtype=torch.long) + lab = torch.zeros((len(batch), mx), dtype=torch.long) + am = torch.zeros((len(batch), mx), dtype=torch.long) + for i, (x, y) in enumerate(batch): + ids[i, : len(x)] = torch.tensor(x) + lab[i, : len(y)] = torch.tensor(y) + am[i, : len(x)] = 1 + return ids, am, lab + + val_rows = units[-N_VAL:] + train_rows = units[:-N_VAL] + train_loader = DataLoader( + PlaneDS(train_rows), batch_size=16, shuffle=True, collate_fn=collate, + num_workers=4, drop_last=True, + ) + total_steps = len(train_loader) * EPOCHS + print(f"[data] train={len(train_rows)} val={len(val_rows)} steps={total_steps} " + f"classes={n_classes}", flush=True) + + backbone = AutoModelForSeq2SeqLM.from_pretrained("google/byt5-small") + encoder = backbone.encoder.cuda().train() # pretrained; decoder discarded + d_model = encoder.config.d_model + + plane_emb = nn.Embedding(n_classes + 1, d_model).cuda().train() + nn.init.normal_(plane_emb.weight, std=0.02) + head = nn.Linear(d_model, n_classes).cuda().train() + + params = list(encoder.parameters()) + list(plane_emb.parameters()) + list(head.parameters()) + optimizer = torch.optim.AdamW(params, lr=1e-4, weight_decay=0.01) + scheduler = get_cosine_schedule_with_warmup(optimizer, 500, total_steps) + ce = nn.CrossEntropyLoss() + + def usable(ck: Path) -> bool: + m = ck / "labels.sha" + return m.exists() and m.read_text().strip() == labels_digest + + ckpts = sorted((c for c in ck_dir.glob("step-*") if usable(c)), + key=lambda p: int(p.name.split("-")[1])) + step = 0 + if ckpts: + sd = torch.load(ckpts[-1] / "model.pt", map_location="cpu", weights_only=True) + encoder.load_state_dict(sd["encoder"]) + plane_emb.load_state_dict(sd["plane_emb"]) + head.load_state_dict(sd["head"]) + optimizer.load_state_dict(torch.load(ckpts[-1] / "optim.pt", map_location="cpu", weights_only=True)) + step = int(ckpts[-1].name.split("-")[1]) + for _ in range(step): + scheduler.step() + print(f"[resume] step-{step}", flush=True) + + def forward_logits(ids, am, plane_in): + embeds = encoder.embed_tokens(ids) + plane_emb(plane_in) + out = encoder(inputs_embeds=embeds, attention_mask=am) + return head(out.last_hidden_state) + + for _epoch in range(EPOCHS): + for ids, am, lab in train_loader: + if step >= total_steps: + break + ids, am, lab = ids.cuda(), am.cuda(), lab.cuda() + q = random.random() # corruption level; q near 1 = first pass + keep = (torch.rand_like(lab, dtype=torch.float) >= q).long() * am + plane_in = torch.where(keep == 1, lab, torch.full_like(lab, MASK_ID)) + with torch.autocast("cuda", torch.bfloat16): + logits = forward_logits(ids, am, plane_in) + loss = ce(logits.float().reshape(-1, n_classes), lab.reshape(-1)) + loss.backward() + torch.nn.utils.clip_grad_norm_(params, 1.0) + optimizer.step() + scheduler.step() + optimizer.zero_grad() + step += 1 + if step % 50 == 0: + print(f"[{RUN} step {step}/{total_steps}] ce={float(loss):.4f} q={q:.2f}", flush=True) + if step % 500 == 0: + ck = ck_dir / f"step-{step}" + ck.mkdir(exist_ok=True) + (ck / "labels.sha").write_text(labels_digest) + torch.save({"encoder": encoder.state_dict(), + "plane_emb": plane_emb.state_dict(), + "head": head.state_dict()}, ck / "model.pt") + torch.save(optimizer.state_dict(), ck / "optim.pt") + checkpoints_volume.commit() + print(f"[volume] committed at step {step}", flush=True) + + best = ck_dir / "best" + best.mkdir(exist_ok=True) + torch.save({"encoder": encoder.state_dict(), + "plane_emb": plane_emb.state_dict(), + "head": head.state_dict(), + "classes": classes}, best / "model.pt") + checkpoints_volume.commit() + + # ---- ID gate: windowed zero-skip SadeedDiac, K-pass ---- + import pandas as pd + import pyarrow.parquet as pq + + encoder.eval() + plane_emb.eval() + head.eval() + table = pq.read_table("data/sadeed-diac-25/train.parquet") + inputs = [DIACRITICS_RE.sub("", t) for t in table.column("input").to_pylist()] + outputs = table.column("output").to_pylist() + + all_windows: list[str] = [] + counts: list[int] = [] + for text in inputs: + ws = split_windows(text) + counts.append(len(ws)) + all_windows.extend(ws) + print(f"[eval] {len(inputs)} paragraphs -> {len(all_windows)} windows", flush=True) + + prog = ck_dir / "eval_progress.jsonl" + saved: dict[int, str] = {} + if prog.exists(): + for line in prog.read_text(encoding="utf-8").splitlines(): + if line.strip(): + row = json.loads(line) + saved[row["i"]] = row["pred"] + print(f"[gen] resuming with {len(saved)} saved windows", flush=True) + + @torch.no_grad() + def predict_window(skel: str) -> str: + ids, _ = encode_unit(skel) + t_ids = torch.tensor([ids], device="cuda") + am = torch.ones_like(t_ids) + plane = torch.full_like(t_ids, MASK_ID) + preds = None + for _ in range(K_PASSES): + with torch.autocast("cuda", torch.bfloat16): + logits = forward_logits(t_ids, am, plane) + preds = logits.argmax(-1) + plane = preds + # per-char class = majority vote of the char's byte-token + # predictions (multi-byte chars span >1 token) + pos = 0 + char_preds: list[str] = [] + for ch in skel: + n = len(ch.encode("utf-8")) + votes = preds[0, pos : pos + n].tolist() + cid = max(set(votes), key=votes.count) + char_preds.append(classes[cid] if cid < n_classes else "") + pos += n + return HP.render(skel, char_preds) + + missing = [i for i in range(len(all_windows)) if i not in saved] + n_new = 0 + with prog.open("a", encoding="utf-8") as prog_out: + for bi in range(0, len(missing), 8): + idxs = missing[bi : bi + 8] + for i in idxs: + pred = predict_window(all_windows[i]) + prog_out.write(json.dumps({"i": i, "pred": pred}, ensure_ascii=False) + "\n") + saved[i] = pred + n_new += len(idxs) + if n_new % 160 == 0: + prog_out.flush() + checkpoints_volume.commit() + print(f"[gen] {len(saved)}/{len(all_windows)} (committed)", flush=True) + checkpoints_volume.commit() + + k = 0 + paragraphs = [] + for _text, c in zip(inputs, counts): + paragraphs.append(" ".join(saved[i] for i in range(k, k + c))) + k += c + + csv_path = Path("/tmp/sadeed_plane_windowed.csv") + pd.DataFrame({"gt": outputs, "pred": paragraphs}).to_csv(csv_path, index=False, header=False) + (ck_dir / "sadeed_preds_windowed.csv").write_text(csv_path.read_text(), encoding="utf-8") + checkpoints_volume.commit() + + from sadeed_evaluator import ArabicDiacritizationEvaluator as E + + print("\n===== run-017 plane, windowed zero-skip (ID gate; " + "rung gate 4.5701 / SOTA gate 2.2864) =====", 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) + + done_marker.touch() + checkpoints_volume.commit() + return {"run": RUN, "steps": step, "k_passes": K_PASSES} + + +@app.local_entrypoint() +def main(): + handle = train.spawn() + print(f"spawned {handle.object_id}; completion = EVAL_DONE at " + f"rababa_checkpoints:{RUN}/EVAL_DONE", flush=True) From 4a328667f2f6a7022852697610f0294323e85436 Mon Sep 17 00:00:00 2001 From: Ronald Tse Date: Sun, 4 Oct 2026 05:06:58 +0800 Subject: [PATCH 03/12] feat(run-019): plane teacher-tier - ByT5-LARGE encoder, 551k units, 1 epoch, K=2 - SOTA gate 2.2864 --- train_arabic_plane_large.py | 371 ++++++++++++++++++++++++++++++++++++ 1 file changed, 371 insertions(+) create mode 100644 train_arabic_plane_large.py diff --git a/train_arabic_plane_large.py b/train_arabic_plane_large.py new file mode 100644 index 0000000..f79f88e --- /dev/null +++ b/train_arabic_plane_large.py @@ -0,0 +1,371 @@ +"""Arabic run-017 — plane-factorized diacritization (Stoicheia WO). + +TODO.sota-2026/04, entered 2026-10-03 (owner: "make it work - +ultimately the best SOTA"). The task is a LENGTH-PRESERVING PLANE +PROJECTION: letters plane given (skeleton), haraqat plane predicted +per position with BIDIRECTIONAL context - the exact component the +left-to-right seq2seq rungs cannot see (iʿrāb depends on sentence +structure ahead; 33% of the residual is word-final case endings). + +Model: ByT5-small ENCODER (pretrained; byte+3 table reused) + a +diacritic-plane embedding (n_classes + MASK) + a per-position +classification head over the corpus haraqat-combo inventory. +Mask-Predict self-conditioning: trained at random corruption levels +q ~ U(0,1]; inference runs K passes of fully parallel prediction +(pass k>1 conditions on pass k-1 argmax). + +Corpus: benchmark-convention RUNNING TEXT only (r5-units + news mix) +- the r9 convention-drift lesson applied by design. NO paradigm tables. + +Gates (pre-registered, windowed zero-skip SadeedDiac-25 full 1,200): +- rung gate: beat the student tier 4.5701 (on-device class); +- SOTA-dedicated gate: beat the teacher tier 2.2864; +- CPU latency benchmark vs ara-diac-small-int8static-2.1 (post-hoc). + +Usage: + modal run --detach train_arabic_plane.py +""" + +from __future__ import annotations + +import random +import re +from pathlib import Path + +import modal + +datasets_volume = modal.Volume.from_name("rababa-datasets", create_if_missing=True) +checkpoints_volume = modal.Volume.from_name("rababa-checkpoints", create_if_missing=True) + +RUN = "rababa_arabic_plane/run-019-plane-large" +UNIT_BYTES = 1400 +N_VAL = 2_000 +EPOCHS = 1 +# teacher-tier: ByT5-LARGE encoder (~1.3B) - SOTA-dedicated gate 2.2864 +MAX_LINES = 600_000 +NEWS_UPSAMPLE = 3 +GOLD_UPSAMPLE = 4 +K_PASSES = 2 + +DIACRITICS_RE = re.compile("[ؐ-ًؚ-ٰٟۖ-ۜ۟-۪ۨ-ۭ]") + +image = ( + modal.Image.debian_slim(python_version="3.11") + .pip_install( + "torch==2.5.1", + "transformers==4.46.3", + "accelerate>=1.1.0", + "pandas", + "tqdm", + "pyarrow", + "pyarabic", + "prettytable", + ) + .add_local_file("sadeed_evaluator.py", "/opt/rababa/sadeed_evaluator.py", copy=True) + .add_local_file("haraqat_planes.py", "/opt/rababa/haraqat_planes.py", copy=True) + .add_local_dir("data/sadeed-diac-25", "/opt/rababa/data/sadeed-diac-25", copy=True) + .workdir("/opt/rababa") + .env({"PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True"}) +) + +app = modal.App("rababa-arabic-plane-large", image=image) + + +def make_unit(text: str) -> str | None: + t = text.strip() + if not t or DIACRITICS_RE.sub("", t).encode("utf-8") == b"": + return None + if len(t.encode("utf-8")) > 1450: + return None + return t + + +def split_windows(text: str, budget: int = UNIT_BYTES) -> list[str]: + if len(text.encode("utf-8")) <= budget: + return [text] + words = text.split() + wins, cur, n = [], [], 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 + + +@app.function( + gpu="A100-80GB", + timeout=24 * 60 * 60, + volumes={"/datasets": datasets_volume, "/checkpoints": checkpoints_volume}, +) +def train() -> dict: + import hashlib + import json + + import torch + import torch.nn as nn + from torch.utils.data import DataLoader, Dataset + from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, get_cosine_schedule_with_warmup + + import haraqat_planes as HP + + datasets_volume.reload() + checkpoints_volume.reload() + + done_marker = Path("/checkpoints") / RUN / "EVAL_DONE" + if done_marker.exists(): + return {"run": RUN, "status": "already-done"} + + # ---- corpus: benchmark-convention running text ---- + cache = Path("/datasets/r5-units") + news_dir = Path("/datasets/arabic-news-r5") + lines: list[str] = [] + for name in ("domain.txt", "replay.txt"): + lines += [l for l in (cache / name).read_text(encoding="utf-8").splitlines() if l.strip()] + news = [l for l in (news_dir / "news.txt").read_text(encoding="utf-8").splitlines() if l.strip()] + gold = [l for l in (news_dir / "wikinews2014_gold.txt").read_text(encoding="utf-8").splitlines() if l.strip()] + units = lines + news * NEWS_UPSAMPLE + gold * GOLD_UPSAMPLE + random.Random(42).shuffle(units) + units = [u for u in (make_unit(u) for u in units[:MAX_LINES]) if u] + print(f"[data] units={len(units)}", flush=True) + + # ---- inventory (canonical combos over the corpus) ---- + ck_dir = Path("/checkpoints") / RUN + ck_dir.mkdir(parents=True, exist_ok=True) + inv_path = ck_dir / "inventory.json" + if inv_path.exists(): + classes = json.loads(inv_path.read_text(encoding="utf-8")) + else: + classes, counts = HP.plane_classes(units) + inv_path.write_text(json.dumps(classes, ensure_ascii=False), encoding="utf-8") + print(f"[inventory] {len(classes)} classes (incl. none); top: " + f"{[(c, counts[c]) for c in classes[1:9]]}", flush=True) + class_to_id = {c: i for i, c in enumerate(classes)} + n_classes = len(classes) + MASK_ID = n_classes # plane vocab = n_classes + 1 (mask) + labels_digest = hashlib.sha256(inv_path.read_bytes()).hexdigest()[:12] + checkpoints_volume.commit() + + tok = AutoTokenizer.from_pretrained("google/byt5-large") + + def encode_unit(text: str): + skel, labels = HP.split_planes(text) + ids = tok(skel).input_ids[:-1] # drop EOS: one token per byte + # char -> byte-span map (Arabic chars are multi-byte) + label_ids = [0] * len(ids) + pos = 0 + for ch, lab in zip(skel, labels): + n = len(ch.encode("utf-8")) + lab = lab.split("\x00", 1)[-1] + cid = class_to_id.get(HP.canon_combo(lab), 0) + for t in range(pos, min(pos + n, len(label_ids))): + label_ids[t] = cid + pos += n + return ids, label_ids + + class PlaneDS(Dataset): + def __init__(self, rows): + self.rows = rows + + def __len__(self): + return len(self.rows) + + def __getitem__(self, i): + return encode_unit(self.rows[i]) + + def collate(batch): + mx = max(len(x[0]) for x in batch) + ids = torch.zeros((len(batch), mx), dtype=torch.long) + lab = torch.zeros((len(batch), mx), dtype=torch.long) + am = torch.zeros((len(batch), mx), dtype=torch.long) + for i, (x, y) in enumerate(batch): + ids[i, : len(x)] = torch.tensor(x) + lab[i, : len(y)] = torch.tensor(y) + am[i, : len(x)] = 1 + return ids, am, lab + + val_rows = units[-N_VAL:] + train_rows = units[:-N_VAL] + train_loader = DataLoader( + PlaneDS(train_rows), batch_size=8, shuffle=True, collate_fn=collate, + num_workers=4, drop_last=True, + ) + total_steps = len(train_loader) * EPOCHS + print(f"[data] train={len(train_rows)} val={len(val_rows)} steps={total_steps} " + f"classes={n_classes}", flush=True) + + backbone = AutoModelForSeq2SeqLM.from_pretrained("google/byt5-large") + encoder = backbone.encoder.cuda().train() # pretrained; decoder discarded + d_model = encoder.config.d_model + + plane_emb = nn.Embedding(n_classes + 1, d_model).cuda().train() + nn.init.normal_(plane_emb.weight, std=0.02) + head = nn.Linear(d_model, n_classes).cuda().train() + + params = list(encoder.parameters()) + list(plane_emb.parameters()) + list(head.parameters()) + optimizer = torch.optim.AdamW(params, lr=1e-4, weight_decay=0.01) + scheduler = get_cosine_schedule_with_warmup(optimizer, 500, total_steps) + ce = nn.CrossEntropyLoss() + + def usable(ck: Path) -> bool: + m = ck / "labels.sha" + return m.exists() and m.read_text().strip() == labels_digest + + ckpts = sorted((c for c in ck_dir.glob("step-*") if usable(c)), + key=lambda p: int(p.name.split("-")[1])) + step = 0 + if ckpts: + sd = torch.load(ckpts[-1] / "model.pt", map_location="cpu", weights_only=True) + encoder.load_state_dict(sd["encoder"]) + plane_emb.load_state_dict(sd["plane_emb"]) + head.load_state_dict(sd["head"]) + optimizer.load_state_dict(torch.load(ckpts[-1] / "optim.pt", map_location="cpu", weights_only=True)) + step = int(ckpts[-1].name.split("-")[1]) + for _ in range(step): + scheduler.step() + print(f"[resume] step-{step}", flush=True) + + def forward_logits(ids, am, plane_in): + embeds = encoder.embed_tokens(ids) + plane_emb(plane_in) + out = encoder(inputs_embeds=embeds, attention_mask=am) + return head(out.last_hidden_state) + + for _epoch in range(EPOCHS): + for ids, am, lab in train_loader: + if step >= total_steps: + break + ids, am, lab = ids.cuda(), am.cuda(), lab.cuda() + q = random.random() # corruption level; q near 1 = first pass + keep = (torch.rand_like(lab, dtype=torch.float) >= q).long() * am + plane_in = torch.where(keep == 1, lab, torch.full_like(lab, MASK_ID)) + with torch.autocast("cuda", torch.bfloat16): + logits = forward_logits(ids, am, plane_in) + loss = ce(logits.float().reshape(-1, n_classes), lab.reshape(-1)) + loss.backward() + torch.nn.utils.clip_grad_norm_(params, 1.0) + optimizer.step() + scheduler.step() + optimizer.zero_grad() + step += 1 + if step % 50 == 0: + print(f"[{RUN} step {step}/{total_steps}] ce={float(loss):.4f} q={q:.2f}", flush=True) + if step % 500 == 0: + ck = ck_dir / f"step-{step}" + ck.mkdir(exist_ok=True) + (ck / "labels.sha").write_text(labels_digest) + torch.save({"encoder": encoder.state_dict(), + "plane_emb": plane_emb.state_dict(), + "head": head.state_dict()}, ck / "model.pt") + torch.save(optimizer.state_dict(), ck / "optim.pt") + checkpoints_volume.commit() + print(f"[volume] committed at step {step}", flush=True) + + best = ck_dir / "best" + best.mkdir(exist_ok=True) + torch.save({"encoder": encoder.state_dict(), + "plane_emb": plane_emb.state_dict(), + "head": head.state_dict(), + "classes": classes}, best / "model.pt") + checkpoints_volume.commit() + + # ---- ID gate: windowed zero-skip SadeedDiac, K-pass ---- + import pandas as pd + import pyarrow.parquet as pq + + encoder.eval() + plane_emb.eval() + head.eval() + table = pq.read_table("data/sadeed-diac-25/train.parquet") + inputs = [DIACRITICS_RE.sub("", t) for t in table.column("input").to_pylist()] + outputs = table.column("output").to_pylist() + + all_windows: list[str] = [] + counts: list[int] = [] + for text in inputs: + ws = split_windows(text) + counts.append(len(ws)) + all_windows.extend(ws) + print(f"[eval] {len(inputs)} paragraphs -> {len(all_windows)} windows", flush=True) + + prog = ck_dir / "eval_progress.jsonl" + saved: dict[int, str] = {} + if prog.exists(): + for line in prog.read_text(encoding="utf-8").splitlines(): + if line.strip(): + row = json.loads(line) + saved[row["i"]] = row["pred"] + print(f"[gen] resuming with {len(saved)} saved windows", flush=True) + + @torch.no_grad() + def predict_window(skel: str) -> str: + ids, _ = encode_unit(skel) + t_ids = torch.tensor([ids], device="cuda") + am = torch.ones_like(t_ids) + plane = torch.full_like(t_ids, MASK_ID) + preds = None + for _ in range(K_PASSES): + with torch.autocast("cuda", torch.bfloat16): + logits = forward_logits(t_ids, am, plane) + preds = logits.argmax(-1) + plane = preds + # per-char class = majority vote of the char's byte-token + # predictions (multi-byte chars span >1 token) + pos = 0 + char_preds: list[str] = [] + for ch in skel: + n = len(ch.encode("utf-8")) + votes = preds[0, pos : pos + n].tolist() + cid = max(set(votes), key=votes.count) + char_preds.append(classes[cid] if cid < n_classes else "") + pos += n + return HP.render(skel, char_preds) + + missing = [i for i in range(len(all_windows)) if i not in saved] + n_new = 0 + with prog.open("a", encoding="utf-8") as prog_out: + for bi in range(0, len(missing), 8): + idxs = missing[bi : bi + 8] + for i in idxs: + pred = predict_window(all_windows[i]) + prog_out.write(json.dumps({"i": i, "pred": pred}, ensure_ascii=False) + "\n") + saved[i] = pred + n_new += len(idxs) + if n_new % 160 == 0: + prog_out.flush() + checkpoints_volume.commit() + print(f"[gen] {len(saved)}/{len(all_windows)} (committed)", flush=True) + checkpoints_volume.commit() + + k = 0 + paragraphs = [] + for _text, c in zip(inputs, counts): + paragraphs.append(" ".join(saved[i] for i in range(k, k + c))) + k += c + + csv_path = Path("/tmp/sadeed_plane_windowed.csv") + pd.DataFrame({"gt": outputs, "pred": paragraphs}).to_csv(csv_path, index=False, header=False) + (ck_dir / "sadeed_preds_windowed.csv").write_text(csv_path.read_text(), encoding="utf-8") + checkpoints_volume.commit() + + from sadeed_evaluator import ArabicDiacritizationEvaluator as E + + print("\n===== run-017 plane, windowed zero-skip (ID gate; " + "rung gate 4.5701 / SOTA gate 2.2864) =====", 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) + + done_marker.touch() + checkpoints_volume.commit() + return {"run": RUN, "steps": step, "k_passes": K_PASSES} + + +@app.local_entrypoint() +def main(): + handle = train.spawn() + print(f"spawned {handle.object_id}; completion = EVAL_DONE at " + f"rababa_checkpoints:{RUN}/EVAL_DONE", flush=True) From 552a6906f32b20021b34c1fb9d7ec38fb80e15fa Mon Sep 17 00:00:00 2001 From: Ronald Tse Date: Sun, 4 Oct 2026 23:29:53 +0800 Subject: [PATCH 04/12] feat(plane arms): run-019 epoch-2 extension + run-020 r7-label distillation run-019 measured 3.0078 at 1 epoch (beats Gemini-Flash, gate open). Arm A: EPOCHS=2 resume (epoch effect measured -1.17pp at small tier). Arm B: ByT5-small plane + r7 teacher labels x6, 3 epochs - the supervision lever, single-variable vs run-018's raw text. --- train_arabic_plane_distill.py | 375 +++++++++++++++++++++++++++++++++ train_arabic_plane_large_e2.py | 371 ++++++++++++++++++++++++++++++++ 2 files changed, 746 insertions(+) create mode 100644 train_arabic_plane_distill.py create mode 100644 train_arabic_plane_large_e2.py diff --git a/train_arabic_plane_distill.py b/train_arabic_plane_distill.py new file mode 100644 index 0000000..c8370c9 --- /dev/null +++ b/train_arabic_plane_distill.py @@ -0,0 +1,375 @@ +"""Arabic run-017 — plane-factorized diacritization (Stoicheia WO). + +TODO.sota-2026/04, entered 2026-10-03 (owner: "make it work - +ultimately the best SOTA"). The task is a LENGTH-PRESERVING PLANE +PROJECTION: letters plane given (skeleton), haraqat plane predicted +per position with BIDIRECTIONAL context - the exact component the +left-to-right seq2seq rungs cannot see (iʿrāb depends on sentence +structure ahead; 33% of the residual is word-final case endings). + +Model: ByT5-small ENCODER (pretrained; byte+3 table reused) + a +diacritic-plane embedding (n_classes + MASK) + a per-position +classification head over the corpus haraqat-combo inventory. +Mask-Predict self-conditioning: trained at random corruption levels +q ~ U(0,1]; inference runs K passes of fully parallel prediction +(pass k>1 conditions on pass k-1 argmax). + +Corpus: benchmark-convention RUNNING TEXT only (r5-units + news mix) +- the r9 convention-drift lesson applied by design. NO paradigm tables. + +Gates (pre-registered, windowed zero-skip SadeedDiac-25 full 1,200): +- rung gate: beat the student tier 4.5701 (on-device class); +- SOTA-dedicated gate: beat the teacher tier 2.2864; +- CPU latency benchmark vs ara-diac-small-int8static-2.1 (post-hoc). + +Usage: + modal run --detach train_arabic_plane.py +""" + +from __future__ import annotations + +import random +import re +from pathlib import Path + +import modal + +datasets_volume = modal.Volume.from_name("rababa-datasets", create_if_missing=True) +checkpoints_volume = modal.Volume.from_name("rababa-checkpoints", create_if_missing=True) + +RUN = "rababa_arabic_plane/run-020-plane-distill" +UNIT_BYTES = 1400 +N_VAL = 2_000 +EPOCHS = 3 +MAX_LINES = 600_000 +NEWS_UPSAMPLE = 3 +GOLD_UPSAMPLE = 4 +K_PASSES = 2 + +DIACRITICS_RE = re.compile("[ؐ-ًؚ-ٰٟۖ-ۜ۟-۪ۨ-ۭ]") + +image = ( + modal.Image.debian_slim(python_version="3.11") + .pip_install( + "torch==2.5.1", + "transformers==4.46.3", + "accelerate>=1.1.0", + "pandas", + "tqdm", + "pyarrow", + "pyarabic", + "prettytable", + ) + .add_local_file("sadeed_evaluator.py", "/opt/rababa/sadeed_evaluator.py", copy=True) + .add_local_file("haraqat_planes.py", "/opt/rababa/haraqat_planes.py", copy=True) + .add_local_dir("data/sadeed-diac-25", "/opt/rababa/data/sadeed-diac-25", copy=True) + .workdir("/opt/rababa") + .env({"PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True"}) +) + +app = modal.App("rababa-arabic-plane-distill", image=image) + + +def make_unit(text: str) -> str | None: + t = text.strip() + if not t or DIACRITICS_RE.sub("", t).encode("utf-8") == b"": + return None + if len(t.encode("utf-8")) > 1450: + return None + return t + + +def split_windows(text: str, budget: int = UNIT_BYTES) -> list[str]: + if len(text.encode("utf-8")) <= budget: + return [text] + words = text.split() + wins, cur, n = [], [], 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 + + +@app.function( + gpu="A100-80GB", + timeout=24 * 60 * 60, + volumes={"/datasets": datasets_volume, "/checkpoints": checkpoints_volume}, +) +def train() -> dict: + import hashlib + import json + + import torch + import torch.nn as nn + from torch.utils.data import DataLoader, Dataset + from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, get_cosine_schedule_with_warmup + + import haraqat_planes as HP + + datasets_volume.reload() + checkpoints_volume.reload() + + done_marker = Path("/checkpoints") / RUN / "EVAL_DONE" + if done_marker.exists(): + return {"run": RUN, "status": "already-done"} + + # ---- corpus: benchmark-convention running text ---- + # distillation: (src skeleton -> r7 teacher's diacritization) pairs + import json as _json + labels_path = Path("/checkpoints/rababa_arabic_distill_small/run-007-r7-muon-6ep/teacher_labels_r7.jsonl") + units = [] + for line in labels_path.read_text(encoding="utf-8").splitlines(): + if not line.strip(): + continue + row = _json.loads(line) + tgt = row["teacher"] + u = make_unit(tgt) + if u: + units.append(u) + # 29,322 label rows: upsample x6 to match run-018's per-epoch scale + units = units * 6 + random.Random(42).shuffle(units) + print(f"[data] distill units={len(units)} (r7 teacher labels x6)", flush=True) + + # ---- inventory (canonical combos over the corpus) ---- + ck_dir = Path("/checkpoints") / RUN + ck_dir.mkdir(parents=True, exist_ok=True) + inv_path = ck_dir / "inventory.json" + if inv_path.exists(): + classes = json.loads(inv_path.read_text(encoding="utf-8")) + else: + classes, counts = HP.plane_classes(units) + inv_path.write_text(json.dumps(classes, ensure_ascii=False), encoding="utf-8") + print(f"[inventory] {len(classes)} classes (incl. none); top: " + f"{[(c, counts[c]) for c in classes[1:9]]}", flush=True) + class_to_id = {c: i for i, c in enumerate(classes)} + n_classes = len(classes) + MASK_ID = n_classes # plane vocab = n_classes + 1 (mask) + labels_digest = hashlib.sha256(inv_path.read_bytes()).hexdigest()[:12] + checkpoints_volume.commit() + + tok = AutoTokenizer.from_pretrained("google/byt5-small") + + def encode_unit(text: str): + skel, labels = HP.split_planes(text) + ids = tok(skel).input_ids[:-1] # drop EOS: one token per byte + # char -> byte-span map (Arabic chars are multi-byte) + label_ids = [0] * len(ids) + pos = 0 + for ch, lab in zip(skel, labels): + n = len(ch.encode("utf-8")) + lab = lab.split("\x00", 1)[-1] + cid = class_to_id.get(HP.canon_combo(lab), 0) + for t in range(pos, min(pos + n, len(label_ids))): + label_ids[t] = cid + pos += n + return ids, label_ids + + class PlaneDS(Dataset): + def __init__(self, rows): + self.rows = rows + + def __len__(self): + return len(self.rows) + + def __getitem__(self, i): + return encode_unit(self.rows[i]) + + def collate(batch): + mx = max(len(x[0]) for x in batch) + ids = torch.zeros((len(batch), mx), dtype=torch.long) + lab = torch.zeros((len(batch), mx), dtype=torch.long) + am = torch.zeros((len(batch), mx), dtype=torch.long) + for i, (x, y) in enumerate(batch): + ids[i, : len(x)] = torch.tensor(x) + lab[i, : len(y)] = torch.tensor(y) + am[i, : len(x)] = 1 + return ids, am, lab + + val_rows = units[-N_VAL:] + train_rows = units[:-N_VAL] + train_loader = DataLoader( + PlaneDS(train_rows), batch_size=16, shuffle=True, collate_fn=collate, + num_workers=4, drop_last=True, + ) + total_steps = len(train_loader) * EPOCHS + print(f"[data] train={len(train_rows)} val={len(val_rows)} steps={total_steps} " + f"classes={n_classes}", flush=True) + + backbone = AutoModelForSeq2SeqLM.from_pretrained("google/byt5-small") + encoder = backbone.encoder.cuda().train() # pretrained; decoder discarded + d_model = encoder.config.d_model + + plane_emb = nn.Embedding(n_classes + 1, d_model).cuda().train() + nn.init.normal_(plane_emb.weight, std=0.02) + head = nn.Linear(d_model, n_classes).cuda().train() + + params = list(encoder.parameters()) + list(plane_emb.parameters()) + list(head.parameters()) + optimizer = torch.optim.AdamW(params, lr=1e-4, weight_decay=0.01) + scheduler = get_cosine_schedule_with_warmup(optimizer, 500, total_steps) + ce = nn.CrossEntropyLoss() + + def usable(ck: Path) -> bool: + m = ck / "labels.sha" + return m.exists() and m.read_text().strip() == labels_digest + + ckpts = sorted((c for c in ck_dir.glob("step-*") if usable(c)), + key=lambda p: int(p.name.split("-")[1])) + step = 0 + if ckpts: + sd = torch.load(ckpts[-1] / "model.pt", map_location="cpu", weights_only=True) + encoder.load_state_dict(sd["encoder"]) + plane_emb.load_state_dict(sd["plane_emb"]) + head.load_state_dict(sd["head"]) + optimizer.load_state_dict(torch.load(ckpts[-1] / "optim.pt", map_location="cpu", weights_only=True)) + step = int(ckpts[-1].name.split("-")[1]) + for _ in range(step): + scheduler.step() + print(f"[resume] step-{step}", flush=True) + + def forward_logits(ids, am, plane_in): + embeds = encoder.embed_tokens(ids) + plane_emb(plane_in) + out = encoder(inputs_embeds=embeds, attention_mask=am) + return head(out.last_hidden_state) + + for _epoch in range(EPOCHS): + for ids, am, lab in train_loader: + if step >= total_steps: + break + ids, am, lab = ids.cuda(), am.cuda(), lab.cuda() + q = random.random() # corruption level; q near 1 = first pass + keep = (torch.rand_like(lab, dtype=torch.float) >= q).long() * am + plane_in = torch.where(keep == 1, lab, torch.full_like(lab, MASK_ID)) + with torch.autocast("cuda", torch.bfloat16): + logits = forward_logits(ids, am, plane_in) + loss = ce(logits.float().reshape(-1, n_classes), lab.reshape(-1)) + loss.backward() + torch.nn.utils.clip_grad_norm_(params, 1.0) + optimizer.step() + scheduler.step() + optimizer.zero_grad() + step += 1 + if step % 50 == 0: + print(f"[{RUN} step {step}/{total_steps}] ce={float(loss):.4f} q={q:.2f}", flush=True) + if step % 500 == 0: + ck = ck_dir / f"step-{step}" + ck.mkdir(exist_ok=True) + (ck / "labels.sha").write_text(labels_digest) + torch.save({"encoder": encoder.state_dict(), + "plane_emb": plane_emb.state_dict(), + "head": head.state_dict()}, ck / "model.pt") + torch.save(optimizer.state_dict(), ck / "optim.pt") + checkpoints_volume.commit() + print(f"[volume] committed at step {step}", flush=True) + + best = ck_dir / "best" + best.mkdir(exist_ok=True) + torch.save({"encoder": encoder.state_dict(), + "plane_emb": plane_emb.state_dict(), + "head": head.state_dict(), + "classes": classes}, best / "model.pt") + checkpoints_volume.commit() + + # ---- ID gate: windowed zero-skip SadeedDiac, K-pass ---- + import pandas as pd + import pyarrow.parquet as pq + + encoder.eval() + plane_emb.eval() + head.eval() + table = pq.read_table("data/sadeed-diac-25/train.parquet") + inputs = [DIACRITICS_RE.sub("", t) for t in table.column("input").to_pylist()] + outputs = table.column("output").to_pylist() + + all_windows: list[str] = [] + counts: list[int] = [] + for text in inputs: + ws = split_windows(text) + counts.append(len(ws)) + all_windows.extend(ws) + print(f"[eval] {len(inputs)} paragraphs -> {len(all_windows)} windows", flush=True) + + prog = ck_dir / "eval_progress.jsonl" + saved: dict[int, str] = {} + if prog.exists(): + for line in prog.read_text(encoding="utf-8").splitlines(): + if line.strip(): + row = json.loads(line) + saved[row["i"]] = row["pred"] + print(f"[gen] resuming with {len(saved)} saved windows", flush=True) + + @torch.no_grad() + def predict_window(skel: str) -> str: + ids, _ = encode_unit(skel) + t_ids = torch.tensor([ids], device="cuda") + am = torch.ones_like(t_ids) + plane = torch.full_like(t_ids, MASK_ID) + preds = None + for _ in range(K_PASSES): + with torch.autocast("cuda", torch.bfloat16): + logits = forward_logits(t_ids, am, plane) + preds = logits.argmax(-1) + plane = preds + # per-char class = majority vote of the char's byte-token + # predictions (multi-byte chars span >1 token) + pos = 0 + char_preds: list[str] = [] + for ch in skel: + n = len(ch.encode("utf-8")) + votes = preds[0, pos : pos + n].tolist() + cid = max(set(votes), key=votes.count) + char_preds.append(classes[cid] if cid < n_classes else "") + pos += n + return HP.render(skel, char_preds) + + missing = [i for i in range(len(all_windows)) if i not in saved] + n_new = 0 + with prog.open("a", encoding="utf-8") as prog_out: + for bi in range(0, len(missing), 8): + idxs = missing[bi : bi + 8] + for i in idxs: + pred = predict_window(all_windows[i]) + prog_out.write(json.dumps({"i": i, "pred": pred}, ensure_ascii=False) + "\n") + saved[i] = pred + n_new += len(idxs) + if n_new % 160 == 0: + prog_out.flush() + checkpoints_volume.commit() + print(f"[gen] {len(saved)}/{len(all_windows)} (committed)", flush=True) + checkpoints_volume.commit() + + k = 0 + paragraphs = [] + for _text, c in zip(inputs, counts): + paragraphs.append(" ".join(saved[i] for i in range(k, k + c))) + k += c + + csv_path = Path("/tmp/sadeed_plane_windowed.csv") + pd.DataFrame({"gt": outputs, "pred": paragraphs}).to_csv(csv_path, index=False, header=False) + (ck_dir / "sadeed_preds_windowed.csv").write_text(csv_path.read_text(), encoding="utf-8") + checkpoints_volume.commit() + + from sadeed_evaluator import ArabicDiacritizationEvaluator as E + + print("\n===== run-017 plane, windowed zero-skip (ID gate; " + "rung gate 4.5701 / SOTA gate 2.2864) =====", 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) + + done_marker.touch() + checkpoints_volume.commit() + return {"run": RUN, "steps": step, "k_passes": K_PASSES} + + +@app.local_entrypoint() +def main(): + handle = train.spawn() + print(f"spawned {handle.object_id}; completion = EVAL_DONE at " + f"rababa_checkpoints:{RUN}/EVAL_DONE", flush=True) diff --git a/train_arabic_plane_large_e2.py b/train_arabic_plane_large_e2.py new file mode 100644 index 0000000..39f6777 --- /dev/null +++ b/train_arabic_plane_large_e2.py @@ -0,0 +1,371 @@ +"""Arabic run-017 — plane-factorized diacritization (Stoicheia WO). + +TODO.sota-2026/04, entered 2026-10-03 (owner: "make it work - +ultimately the best SOTA"). The task is a LENGTH-PRESERVING PLANE +PROJECTION: letters plane given (skeleton), haraqat plane predicted +per position with BIDIRECTIONAL context - the exact component the +left-to-right seq2seq rungs cannot see (iʿrāb depends on sentence +structure ahead; 33% of the residual is word-final case endings). + +Model: ByT5-small ENCODER (pretrained; byte+3 table reused) + a +diacritic-plane embedding (n_classes + MASK) + a per-position +classification head over the corpus haraqat-combo inventory. +Mask-Predict self-conditioning: trained at random corruption levels +q ~ U(0,1]; inference runs K passes of fully parallel prediction +(pass k>1 conditions on pass k-1 argmax). + +Corpus: benchmark-convention RUNNING TEXT only (r5-units + news mix) +- the r9 convention-drift lesson applied by design. NO paradigm tables. + +Gates (pre-registered, windowed zero-skip SadeedDiac-25 full 1,200): +- rung gate: beat the student tier 4.5701 (on-device class); +- SOTA-dedicated gate: beat the teacher tier 2.2864; +- CPU latency benchmark vs ara-diac-small-int8static-2.1 (post-hoc). + +Usage: + modal run --detach train_arabic_plane.py +""" + +from __future__ import annotations + +import random +import re +from pathlib import Path + +import modal + +datasets_volume = modal.Volume.from_name("rababa-datasets", create_if_missing=True) +checkpoints_volume = modal.Volume.from_name("rababa-checkpoints", create_if_missing=True) + +RUN = "rababa_arabic_plane/run-019-plane-large" +UNIT_BYTES = 1400 +N_VAL = 2_000 +EPOCHS = 2 +# teacher-tier: ByT5-LARGE encoder (~1.3B) - SOTA-dedicated gate 2.2864 +MAX_LINES = 600_000 +NEWS_UPSAMPLE = 3 +GOLD_UPSAMPLE = 4 +K_PASSES = 2 + +DIACRITICS_RE = re.compile("[ؐ-ًؚ-ٰٟۖ-ۜ۟-۪ۨ-ۭ]") + +image = ( + modal.Image.debian_slim(python_version="3.11") + .pip_install( + "torch==2.5.1", + "transformers==4.46.3", + "accelerate>=1.1.0", + "pandas", + "tqdm", + "pyarrow", + "pyarabic", + "prettytable", + ) + .add_local_file("sadeed_evaluator.py", "/opt/rababa/sadeed_evaluator.py", copy=True) + .add_local_file("haraqat_planes.py", "/opt/rababa/haraqat_planes.py", copy=True) + .add_local_dir("data/sadeed-diac-25", "/opt/rababa/data/sadeed-diac-25", copy=True) + .workdir("/opt/rababa") + .env({"PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True"}) +) + +app = modal.App("rababa-arabic-plane-large2", image=image) + + +def make_unit(text: str) -> str | None: + t = text.strip() + if not t or DIACRITICS_RE.sub("", t).encode("utf-8") == b"": + return None + if len(t.encode("utf-8")) > 1450: + return None + return t + + +def split_windows(text: str, budget: int = UNIT_BYTES) -> list[str]: + if len(text.encode("utf-8")) <= budget: + return [text] + words = text.split() + wins, cur, n = [], [], 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 + + +@app.function( + gpu="A100-80GB", + timeout=24 * 60 * 60, + volumes={"/datasets": datasets_volume, "/checkpoints": checkpoints_volume}, +) +def train() -> dict: + import hashlib + import json + + import torch + import torch.nn as nn + from torch.utils.data import DataLoader, Dataset + from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, get_cosine_schedule_with_warmup + + import haraqat_planes as HP + + datasets_volume.reload() + checkpoints_volume.reload() + + done_marker = Path("/checkpoints") / RUN / "EVAL_DONE_E2" # epoch-2 extension: new marker, same run dir (resume from step-*) + if done_marker.exists(): + return {"run": RUN, "status": "already-done"} + + # ---- corpus: benchmark-convention running text ---- + cache = Path("/datasets/r5-units") + news_dir = Path("/datasets/arabic-news-r5") + lines: list[str] = [] + for name in ("domain.txt", "replay.txt"): + lines += [l for l in (cache / name).read_text(encoding="utf-8").splitlines() if l.strip()] + news = [l for l in (news_dir / "news.txt").read_text(encoding="utf-8").splitlines() if l.strip()] + gold = [l for l in (news_dir / "wikinews2014_gold.txt").read_text(encoding="utf-8").splitlines() if l.strip()] + units = lines + news * NEWS_UPSAMPLE + gold * GOLD_UPSAMPLE + random.Random(42).shuffle(units) + units = [u for u in (make_unit(u) for u in units[:MAX_LINES]) if u] + print(f"[data] units={len(units)}", flush=True) + + # ---- inventory (canonical combos over the corpus) ---- + ck_dir = Path("/checkpoints") / RUN + ck_dir.mkdir(parents=True, exist_ok=True) + inv_path = ck_dir / "inventory.json" + if inv_path.exists(): + classes = json.loads(inv_path.read_text(encoding="utf-8")) + else: + classes, counts = HP.plane_classes(units) + inv_path.write_text(json.dumps(classes, ensure_ascii=False), encoding="utf-8") + print(f"[inventory] {len(classes)} classes (incl. none); top: " + f"{[(c, counts[c]) for c in classes[1:9]]}", flush=True) + class_to_id = {c: i for i, c in enumerate(classes)} + n_classes = len(classes) + MASK_ID = n_classes # plane vocab = n_classes + 1 (mask) + labels_digest = hashlib.sha256(inv_path.read_bytes()).hexdigest()[:12] + checkpoints_volume.commit() + + tok = AutoTokenizer.from_pretrained("google/byt5-large") + + def encode_unit(text: str): + skel, labels = HP.split_planes(text) + ids = tok(skel).input_ids[:-1] # drop EOS: one token per byte + # char -> byte-span map (Arabic chars are multi-byte) + label_ids = [0] * len(ids) + pos = 0 + for ch, lab in zip(skel, labels): + n = len(ch.encode("utf-8")) + lab = lab.split("\x00", 1)[-1] + cid = class_to_id.get(HP.canon_combo(lab), 0) + for t in range(pos, min(pos + n, len(label_ids))): + label_ids[t] = cid + pos += n + return ids, label_ids + + class PlaneDS(Dataset): + def __init__(self, rows): + self.rows = rows + + def __len__(self): + return len(self.rows) + + def __getitem__(self, i): + return encode_unit(self.rows[i]) + + def collate(batch): + mx = max(len(x[0]) for x in batch) + ids = torch.zeros((len(batch), mx), dtype=torch.long) + lab = torch.zeros((len(batch), mx), dtype=torch.long) + am = torch.zeros((len(batch), mx), dtype=torch.long) + for i, (x, y) in enumerate(batch): + ids[i, : len(x)] = torch.tensor(x) + lab[i, : len(y)] = torch.tensor(y) + am[i, : len(x)] = 1 + return ids, am, lab + + val_rows = units[-N_VAL:] + train_rows = units[:-N_VAL] + train_loader = DataLoader( + PlaneDS(train_rows), batch_size=8, shuffle=True, collate_fn=collate, + num_workers=4, drop_last=True, + ) + total_steps = len(train_loader) * EPOCHS + print(f"[data] train={len(train_rows)} val={len(val_rows)} steps={total_steps} " + f"classes={n_classes}", flush=True) + + backbone = AutoModelForSeq2SeqLM.from_pretrained("google/byt5-large") + encoder = backbone.encoder.cuda().train() # pretrained; decoder discarded + d_model = encoder.config.d_model + + plane_emb = nn.Embedding(n_classes + 1, d_model).cuda().train() + nn.init.normal_(plane_emb.weight, std=0.02) + head = nn.Linear(d_model, n_classes).cuda().train() + + params = list(encoder.parameters()) + list(plane_emb.parameters()) + list(head.parameters()) + optimizer = torch.optim.AdamW(params, lr=1e-4, weight_decay=0.01) + scheduler = get_cosine_schedule_with_warmup(optimizer, 500, total_steps) + ce = nn.CrossEntropyLoss() + + def usable(ck: Path) -> bool: + m = ck / "labels.sha" + return m.exists() and m.read_text().strip() == labels_digest + + ckpts = sorted((c for c in ck_dir.glob("step-*") if usable(c)), + key=lambda p: int(p.name.split("-")[1])) + step = 0 + if ckpts: + sd = torch.load(ckpts[-1] / "model.pt", map_location="cpu", weights_only=True) + encoder.load_state_dict(sd["encoder"]) + plane_emb.load_state_dict(sd["plane_emb"]) + head.load_state_dict(sd["head"]) + optimizer.load_state_dict(torch.load(ckpts[-1] / "optim.pt", map_location="cpu", weights_only=True)) + step = int(ckpts[-1].name.split("-")[1]) + for _ in range(step): + scheduler.step() + print(f"[resume] step-{step}", flush=True) + + def forward_logits(ids, am, plane_in): + embeds = encoder.embed_tokens(ids) + plane_emb(plane_in) + out = encoder(inputs_embeds=embeds, attention_mask=am) + return head(out.last_hidden_state) + + for _epoch in range(EPOCHS): + for ids, am, lab in train_loader: + if step >= total_steps: + break + ids, am, lab = ids.cuda(), am.cuda(), lab.cuda() + q = random.random() # corruption level; q near 1 = first pass + keep = (torch.rand_like(lab, dtype=torch.float) >= q).long() * am + plane_in = torch.where(keep == 1, lab, torch.full_like(lab, MASK_ID)) + with torch.autocast("cuda", torch.bfloat16): + logits = forward_logits(ids, am, plane_in) + loss = ce(logits.float().reshape(-1, n_classes), lab.reshape(-1)) + loss.backward() + torch.nn.utils.clip_grad_norm_(params, 1.0) + optimizer.step() + scheduler.step() + optimizer.zero_grad() + step += 1 + if step % 50 == 0: + print(f"[{RUN} step {step}/{total_steps}] ce={float(loss):.4f} q={q:.2f}", flush=True) + if step % 500 == 0: + ck = ck_dir / f"step-{step}" + ck.mkdir(exist_ok=True) + (ck / "labels.sha").write_text(labels_digest) + torch.save({"encoder": encoder.state_dict(), + "plane_emb": plane_emb.state_dict(), + "head": head.state_dict()}, ck / "model.pt") + torch.save(optimizer.state_dict(), ck / "optim.pt") + checkpoints_volume.commit() + print(f"[volume] committed at step {step}", flush=True) + + best = ck_dir / "best" + best.mkdir(exist_ok=True) + torch.save({"encoder": encoder.state_dict(), + "plane_emb": plane_emb.state_dict(), + "head": head.state_dict(), + "classes": classes}, best / "model.pt") + checkpoints_volume.commit() + + # ---- ID gate: windowed zero-skip SadeedDiac, K-pass ---- + import pandas as pd + import pyarrow.parquet as pq + + encoder.eval() + plane_emb.eval() + head.eval() + table = pq.read_table("data/sadeed-diac-25/train.parquet") + inputs = [DIACRITICS_RE.sub("", t) for t in table.column("input").to_pylist()] + outputs = table.column("output").to_pylist() + + all_windows: list[str] = [] + counts: list[int] = [] + for text in inputs: + ws = split_windows(text) + counts.append(len(ws)) + all_windows.extend(ws) + print(f"[eval] {len(inputs)} paragraphs -> {len(all_windows)} windows", flush=True) + + prog = ck_dir / "eval_progress.jsonl" + saved: dict[int, str] = {} + if prog.exists(): + for line in prog.read_text(encoding="utf-8").splitlines(): + if line.strip(): + row = json.loads(line) + saved[row["i"]] = row["pred"] + print(f"[gen] resuming with {len(saved)} saved windows", flush=True) + + @torch.no_grad() + def predict_window(skel: str) -> str: + ids, _ = encode_unit(skel) + t_ids = torch.tensor([ids], device="cuda") + am = torch.ones_like(t_ids) + plane = torch.full_like(t_ids, MASK_ID) + preds = None + for _ in range(K_PASSES): + with torch.autocast("cuda", torch.bfloat16): + logits = forward_logits(t_ids, am, plane) + preds = logits.argmax(-1) + plane = preds + # per-char class = majority vote of the char's byte-token + # predictions (multi-byte chars span >1 token) + pos = 0 + char_preds: list[str] = [] + for ch in skel: + n = len(ch.encode("utf-8")) + votes = preds[0, pos : pos + n].tolist() + cid = max(set(votes), key=votes.count) + char_preds.append(classes[cid] if cid < n_classes else "") + pos += n + return HP.render(skel, char_preds) + + missing = [i for i in range(len(all_windows)) if i not in saved] + n_new = 0 + with prog.open("a", encoding="utf-8") as prog_out: + for bi in range(0, len(missing), 8): + idxs = missing[bi : bi + 8] + for i in idxs: + pred = predict_window(all_windows[i]) + prog_out.write(json.dumps({"i": i, "pred": pred}, ensure_ascii=False) + "\n") + saved[i] = pred + n_new += len(idxs) + if n_new % 160 == 0: + prog_out.flush() + checkpoints_volume.commit() + print(f"[gen] {len(saved)}/{len(all_windows)} (committed)", flush=True) + checkpoints_volume.commit() + + k = 0 + paragraphs = [] + for _text, c in zip(inputs, counts): + paragraphs.append(" ".join(saved[i] for i in range(k, k + c))) + k += c + + csv_path = Path("/tmp/sadeed_plane_windowed.csv") + pd.DataFrame({"gt": outputs, "pred": paragraphs}).to_csv(csv_path, index=False, header=False) + (ck_dir / "sadeed_preds_windowed.csv").write_text(csv_path.read_text(), encoding="utf-8") + checkpoints_volume.commit() + + from sadeed_evaluator import ArabicDiacritizationEvaluator as E + + print("\n===== run-017 plane, windowed zero-skip (ID gate; " + "rung gate 4.5701 / SOTA gate 2.2864) =====", 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) + + done_marker.touch() + checkpoints_volume.commit() + return {"run": RUN, "steps": step, "k_passes": K_PASSES} + + +@app.local_entrypoint() +def main(): + handle = train.spawn() + print(f"spawned {handle.object_id}; completion = EVAL_DONE at " + f"rababa_checkpoints:{RUN}/EVAL_DONE", flush=True) From 2d2f4d02aa0f862278428b3206c583c41492365a Mon Sep 17 00:00:00 2001 From: Ronald Tse Date: Mon, 5 Oct 2026 20:57:29 +0800 Subject: [PATCH 05/12] fix(run-019 e2): fresh marker + per-marker eval progress file The first E2 eval reused epoch-1's eval_progress.jsonl: every window was already 'saved', so it regenerated nothing and re-scored epoch-1 predictions (DER bit-identical at 3.0078). EVAL_DONE_E2B + eval_progress_EVAL_DONE_E2B.jsonl force full regeneration. --- train_arabic_plane_large_e2.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/train_arabic_plane_large_e2.py b/train_arabic_plane_large_e2.py index 39f6777..e921c79 100644 --- a/train_arabic_plane_large_e2.py +++ b/train_arabic_plane_large_e2.py @@ -116,7 +116,10 @@ def train() -> dict: datasets_volume.reload() checkpoints_volume.reload() - done_marker = Path("/checkpoints") / RUN / "EVAL_DONE_E2" # epoch-2 extension: new marker, same run dir (resume from step-*) + # E2B: the first E2 eval silently reused epoch-1's eval_progress.jsonl + # (all windows "saved" -> zero regeneration -> re-scored epoch-1 preds). + # Fresh marker + per-marker progress file so each eval round regenerates. + done_marker = Path("/checkpoints") / RUN / "EVAL_DONE_E2B" if done_marker.exists(): return {"run": RUN, "status": "already-done"} @@ -291,7 +294,7 @@ def forward_logits(ids, am, plane_in): all_windows.extend(ws) print(f"[eval] {len(inputs)} paragraphs -> {len(all_windows)} windows", flush=True) - prog = ck_dir / "eval_progress.jsonl" + prog = ck_dir / f"eval_progress_{done_marker.name}.jsonl" saved: dict[int, str] = {} if prog.exists(): for line in prog.read_text(encoding="utf-8").splitlines(): From 45ed3808b1ac83900f0059f114341f5ea0997de6 Mon Sep 17 00:00:00 2001 From: Ronald Tse Date: Mon, 5 Oct 2026 22:38:59 +0800 Subject: [PATCH 06/12] feat(plane): artifact zip packaging - contract builder + parameterized exporter build_plane_zip emits the secryst-py PlaneModel.from_zip contract (metadata kind=plane + member sha256s, plane.onnx, classes.json); exporter gains --checkpoint/--backbone/--prefix/--bench-skip. --- export_plane_onnx.py | 30 ++++++++----- package_plane_zip.py | 77 +++++++++++++++++++++++++++++++++ tests/test_package_plane_zip.py | 54 +++++++++++++++++++++++ 3 files changed, 150 insertions(+), 11 deletions(-) create mode 100644 package_plane_zip.py create mode 100644 tests/test_package_plane_zip.py diff --git a/export_plane_onnx.py b/export_plane_onnx.py index 2209d1f..6e77366 100644 --- a/export_plane_onnx.py +++ b/export_plane_onnx.py @@ -8,6 +8,8 @@ Usage: python3 export_plane_onnx.py # export + quantize + bench python3 export_plane_onnx.py --bench-only + python3 export_plane_onnx.py --checkpoint /tmp/ck.pt --backbone google/byt5-small \ + --prefix plane-int8 --bench-skip """ from __future__ import annotations @@ -38,14 +40,14 @@ def fetch_best() -> Path: return ck -def build_model(sd): +def build_model(sd, backbone_name: str = "google/byt5-small"): import torch import torch.nn as nn from transformers import AutoModelForSeq2SeqLM classes = sd["classes"] n_classes = len(classes) - backbone = AutoModelForSeq2SeqLM.from_pretrained("google/byt5-small") + backbone = AutoModelForSeq2SeqLM.from_pretrained(backbone_name) encoder = backbone.encoder encoder.load_state_dict(sd["encoder"]) plane_emb = nn.Embedding(n_classes + 1, encoder.config.d_model) @@ -69,17 +71,18 @@ def forward(self, input_ids, plane_ids): return m, classes -def export() -> None: +def export(checkpoint: Path | None = None, backbone: str = "google/byt5-small", + prefix: str = "plane-1.0") -> None: import torch - sd = torch.load(fetch_best(), map_location="cpu", weights_only=True) - model, classes = build_model(sd) + sd = torch.load(checkpoint or fetch_best(), map_location="cpu", weights_only=True) + model, classes = build_model(sd, backbone) n_classes = len(classes) tok_ids = torch.tensor([[b + 3 for b in SAMPLE.encode("utf-8")[:1400]]], dtype=torch.long) plane = torch.full_like(tok_ids, n_classes) OUT_DIR.mkdir(parents=True, exist_ok=True) - fp32 = OUT_DIR / "plane-1.0-fp32.onnx" + fp32 = OUT_DIR / f"{prefix}-fp32.onnx" torch.onnx.export( model, (tok_ids, plane), str(fp32), input_names=["input_ids", "plane_ids"], @@ -95,14 +98,14 @@ def export() -> None: from onnxruntime.quantization import QuantType, quantize_dynamic - int8 = OUT_DIR / "plane-1.0-int8.onnx" + int8 = OUT_DIR / f"{prefix}-int8.onnx" quantize_dynamic(str(fp32), str(int8), weight_type=QuantType.QInt8) print(f"quantized {int8} ({int8.stat().st_size / 1e6:.1f} MB)") (OUT_DIR / "classes.json").write_text( __import__("json").dumps(classes, ensure_ascii=False), encoding="utf-8") -def bench() -> None: +def bench(prefix: str = "plane-1.0") -> None: import json import numpy as np @@ -112,7 +115,7 @@ def bench() -> None: n_classes = len(classes) ids = np.array([[b + 3 for b in SAMPLE.encode("utf-8")[:1400]]], dtype=np.int64) - for name in ("plane-1.0-fp32.onnx", "plane-1.0-int8.onnx"): + for name in (f"{prefix}-fp32.onnx", f"{prefix}-int8.onnx"): sess = ort.InferenceSession(str(OUT_DIR / name), providers=["CPUExecutionProvider"]) for k in (1, 2, 4): # warmup @@ -134,7 +137,12 @@ def bench() -> None: if __name__ == "__main__": ap = argparse.ArgumentParser() ap.add_argument("--bench-only", action="store_true") + ap.add_argument("--bench-skip", action="store_true") + ap.add_argument("--checkpoint", type=Path, default=None) + ap.add_argument("--backbone", default="google/byt5-small") + ap.add_argument("--prefix", default="plane-1.0") args = ap.parse_args() if not args.bench_only: - export() - bench() + export(args.checkpoint, args.backbone, args.prefix) + if not args.bench_skip: + bench(args.prefix) diff --git a/package_plane_zip.py b/package_plane_zip.py new file mode 100644 index 0000000..6af4119 --- /dev/null +++ b/package_plane_zip.py @@ -0,0 +1,77 @@ +"""Build a plane-model artifact zip per the IMF plane contract. + +Bundle layout (verified by secryst-py PlaneModel.from_zip): + metadata.yaml id, precision, kind=plane, k_passes, classes, member sha256s + plane.onnx one graph (input_ids, plane_ids) -> per-position class logits + classes.json class id -> harakat combo ("" = bare) + +Usage: + python3 package_plane_zip.py --onnx /tmp/plane-onnx/plane-int8.onnx \ + --classes /tmp/plane-onnx/classes.json \ + --id ara-diac-plane-1.0 --precision int8 --out dist/ara-diac-plane-1.0.zip +""" + +from __future__ import annotations + +import argparse +import hashlib +import io +import json +import zipfile +from pathlib import Path + + +def build_plane_zip(model_id: str, precision: str, onnx: bytes, classes: list[str], + k_passes: int = 2) -> bytes: + classes_bytes = json.dumps(classes, ensure_ascii=False, indent=1).encode("utf-8") + members = { + "plane.onnx": onnx, + "classes.json": classes_bytes, + } + meta = ( + "id: {id}\n" + "kind: plane\n" + "precision: {precision}\n" + "k_passes: {k}\n" + "members:\n" + " plane.onnx: {onnx_sha}\n" + " classes.json: {classes_sha}\n" + ).format( + id=model_id, + precision=precision, + k=k_passes, + onnx_sha=hashlib.sha256(onnx).hexdigest(), + classes_sha=hashlib.sha256(classes_bytes).hexdigest(), + ).encode("utf-8") + + buf = io.BytesIO() + with zipfile.ZipFile(buf, "w", zipfile.ZIP_STORED) as zf: + for name, data in [("metadata.yaml", meta), *members.items()]: + zf.writestr(name, data) + return buf.getvalue() + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--onnx", required=True) + ap.add_argument("--classes", required=True) + ap.add_argument("--id", required=True) + ap.add_argument("--precision", required=True) + ap.add_argument("--out", required=True) + ap.add_argument("--k-passes", type=int, default=2) + args = ap.parse_args() + + data = build_plane_zip( + args.id, args.precision, + Path(args.onnx).read_bytes(), + json.loads(Path(args.classes).read_text(encoding="utf-8")), + k_passes=args.k_passes, + ) + out = Path(args.out) + out.parent.mkdir(parents=True, exist_ok=True) + out.write_bytes(data) + print(f"{out} ({len(data) / 1e6:.1f} MB) sha256={hashlib.sha256(data).hexdigest()}") + + +if __name__ == "__main__": + main() diff --git a/tests/test_package_plane_zip.py b/tests/test_package_plane_zip.py new file mode 100644 index 0000000..a9fec64 --- /dev/null +++ b/tests/test_package_plane_zip.py @@ -0,0 +1,54 @@ +"""Tests for the plane-artifact zip builder (IMF plane contract). + +The built zip must be loadable by the secryst-py PlaneModel.from_zip +contract: metadata.yaml (kind=plane, k_passes, members sha256 map) + +plane.onnx + classes.json, every member sha-verified. +""" + +import io +import json +import sys +import unittest +import zipfile +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) + +from package_plane_zip import build_plane_zip # noqa: E402 + +FAKE_ONNX = b"\x08\x01fake-onnx-bytes" * 64 +CLASSES = ["", "َ", "ُ", "ِ", "ّ", "ْ"] + + +class TestBuildPlaneZip(unittest.TestCase): + def test_zip_members_and_metadata(self): + data = build_plane_zip("ara-diac-plane-1.0", "int8", FAKE_ONNX, CLASSES, k_passes=2) + zf = zipfile.ZipFile(io.BytesIO(data)) + self.assertEqual(set(zf.namelist()), {"metadata.yaml", "plane.onnx", "classes.json"}) + meta = __import__("yaml").safe_load(zf.read("metadata.yaml")) + self.assertEqual(meta["kind"], "plane") + self.assertEqual(meta["id"], "ara-diac-plane-1.0") + self.assertEqual(meta["precision"], "int8") + self.assertEqual(meta["k_passes"], 2) + self.assertEqual(json.loads(zf.read("classes.json")), CLASSES) + + def test_member_sha256_recorded_and_correct(self): + import hashlib + + data = build_plane_zip("ara-diac-plane-1.0", "int8", FAKE_ONNX, CLASSES, k_passes=2) + zf = zipfile.ZipFile(io.BytesIO(data)) + meta = __import__("yaml").safe_load(zf.read("metadata.yaml")) + self.assertEqual(meta["members"]["plane.onnx"], hashlib.sha256(FAKE_ONNX).hexdigest()) + self.assertEqual( + meta["members"]["classes.json"], + hashlib.sha256(zf.read("classes.json")).hexdigest(), + ) + + def test_whole_zip_bytes_are_deterministic(self): + a = build_plane_zip("ara-diac-plane-1.0", "int8", FAKE_ONNX, CLASSES, k_passes=2) + b = build_plane_zip("ara-diac-plane-1.0", "int8", FAKE_ONNX, CLASSES, k_passes=2) + self.assertEqual(a, b) + + +if __name__ == "__main__": + unittest.main() From cfe842ea1ae839eee40361c914aeac465e5dda52 Mon Sep 17 00:00:00 2001 From: Ronald Tse Date: Tue, 6 Oct 2026 00:52:49 +0800 Subject: [PATCH 07/12] fix(modal export): /opt on sys.path for the shared builder import --- export_plane_onnx_modal.py | 96 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 96 insertions(+) create mode 100644 export_plane_onnx_modal.py diff --git a/export_plane_onnx_modal.py b/export_plane_onnx_modal.py new file mode 100644 index 0000000..7e3cd6a --- /dev/null +++ b/export_plane_onnx_modal.py @@ -0,0 +1,96 @@ +"""Export the run-019 plane-large artifact ON MODAL. + +The 5.2 GB fp32 checkpoint tears during host-side volume downloads +(two independent `modal volume get` calls returned different bytes), +so the export runs where the volume mounts natively: torch.load, +ONNX fp32 export, dynamic int8 quantize, then the int8 graph + +classes.json are staged back on the volume with their sha256s. + +Usage: + modal run --detach export_plane_onnx_modal.py +""" + +from __future__ import annotations + +import hashlib +import json +from pathlib import Path + +import modal + +app = modal.App("rababa-plane-export-large") + +RUN = "rababa_arabic_plane/run-019-plane-large" +BACKBONE = "google/byt5-large" +PREFIX = "plane-large-e2" +OUT = Path("/checkpoints") / RUN / "export" + +checkpoints_volume = modal.Volume.from_name("rababa-checkpoints", create_if_missing=True) + +image = ( + modal.Image.debian_slim(python_version="3.11") + .pip_install("torch", "transformers", "onnx", "onnxruntime", "numpy") + .add_local_file(Path(__file__).parent / "export_plane_onnx.py", "/opt/export_plane_onnx.py", copy=True) +) + + +@app.function(image=image, timeout=3600, cpu=8, memory=32768, + volumes={"/checkpoints": checkpoints_volume}) +def export() -> dict: + import sys + sys.path.insert(0, "/opt") + import torch + from export_plane_onnx import build_model + + ck = Path("/checkpoints") / RUN / "best" / "model.pt" + sd = torch.load(ck, map_location="cpu", weights_only=True) + model, classes = build_model(sd, BACKBONE) + n_classes = len(classes) + + tok_ids = torch.tensor([[b + 3 for b in ("و" * 100).encode("utf-8")[:1400]]], dtype=torch.long) + plane = torch.full_like(tok_ids, n_classes) + + fp32 = Path("/tmp") / f"{PREFIX}-fp32.onnx" + torch.onnx.export( + model, (tok_ids, plane), str(fp32), + input_names=["input_ids", "plane_ids"], + output_names=["class_logits"], + dynamic_axes={ + "input_ids": {0: "batch", 1: "seq"}, + "plane_ids": {0: "batch", 1: "seq"}, + "class_logits": {0: "batch", 1: "seq"}, + }, + opset_version=14, + ) + print(f"exported {fp32} ({fp32.stat().st_size / 1e6:.1f} MB)", flush=True) + + from onnxruntime.quantization import QuantType, quantize_dynamic + + int8 = Path("/tmp") / f"{PREFIX}-int8.onnx" + quantize_dynamic(str(fp32), str(int8), weight_type=QuantType.QInt8) + + classes_bytes = json.dumps(classes, ensure_ascii=False, indent=1).encode("utf-8") + OUT.mkdir(parents=True, exist_ok=True) + int8_bytes = int8.read_bytes() + (OUT / f"{PREFIX}-int8.onnx").write_bytes(int8_bytes) + (OUT / "classes.json").write_bytes(classes_bytes) + + def sha(b: bytes) -> str: + return hashlib.sha256(b).hexdigest() + + result = { + "int8_sha256": sha(int8_bytes), + "int8_size": len(int8_bytes), + "classes_sha256": sha(classes_bytes), + } + (OUT / "sha256s.json").write_text(json.dumps(result, indent=1), encoding="utf-8") + checkpoints_volume.commit() + print(json.dumps(result, indent=1), flush=True) + return result + + +@app.local_entrypoint() +def main() -> None: + handle = export.spawn() + print(f"spawned {handle.object_id}; output = " + f"rababa_checkpoints:{RUN}/export/{PREFIX}-int8.onnx", flush=True) From 52732523744f17c07936b70139bd7860d23f9b7f Mon Sep 17 00:00:00 2001 From: Ronald Tse Date: Tue, 6 Oct 2026 00:57:58 +0800 Subject: [PATCH 08/12] fix(modal export): onnxscript dep for torch 2.14 dynamo exporter --- export_plane_onnx_modal.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/export_plane_onnx_modal.py b/export_plane_onnx_modal.py index 7e3cd6a..76c7b44 100644 --- a/export_plane_onnx_modal.py +++ b/export_plane_onnx_modal.py @@ -29,7 +29,7 @@ image = ( modal.Image.debian_slim(python_version="3.11") - .pip_install("torch", "transformers", "onnx", "onnxruntime", "numpy") + .pip_install("torch", "transformers", "onnx", "onnxruntime", "onnxscript", "numpy") .add_local_file(Path(__file__).parent / "export_plane_onnx.py", "/opt/export_plane_onnx.py", copy=True) ) From b95cf1783dcf1abfe09208daa9144205226c71e5 Mon Sep 17 00:00:00 2001 From: Ronald Tse Date: Tue, 6 Oct 2026 01:02:19 +0800 Subject: [PATCH 09/12] fix(modal export): legacy ONNX exporter - dynamo cannot trace T5 attention reshape --- export_plane_onnx_modal.py | 1 + 1 file changed, 1 insertion(+) diff --git a/export_plane_onnx_modal.py b/export_plane_onnx_modal.py index 76c7b44..31fba96 100644 --- a/export_plane_onnx_modal.py +++ b/export_plane_onnx_modal.py @@ -61,6 +61,7 @@ def export() -> dict: "class_logits": {0: "batch", 1: "seq"}, }, opset_version=14, + dynamo=False, ) print(f"exported {fp32} ({fp32.stat().st_size / 1e6:.1f} MB)", flush=True) From 54d37af0f296864c088dd7e4cf83087c8d4cb762 Mon Sep 17 00:00:00 2001 From: Ronald Tse Date: Wed, 7 Oct 2026 10:33:24 +0800 Subject: [PATCH 10/12] feat(hebrew): nikud plane module (TDD 7/7) + run-021 plane trainer The Arabic run-018 recipe applied to Hebrew over the v4 combined corpus; eval on nakdimon test with the exact seq2seq_der protocol (heb-diac-1.1's 16.44 DER gate). Per-marker eval progress file (the run-019 E2 lesson applied by design). --- nikud_planes.py | 77 ++++++++++ test_nikud_planes.py | 60 ++++++++ train_hebrew_plane.py | 338 ++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 475 insertions(+) create mode 100644 nikud_planes.py create mode 100644 test_nikud_planes.py create mode 100644 train_hebrew_plane.py diff --git a/nikud_planes.py b/nikud_planes.py new file mode 100644 index 0000000..af0f039 --- /dev/null +++ b/nikud_planes.py @@ -0,0 +1,77 @@ +"""Hebrew nikud plane decomposition (the haraqat_planes analog). + +Splits labeled Hebrew text into a letter skeleton plus one nikud +cluster per letter. Marks FOLLOW their base letter and close into the +previous letter's class; the cluster is canon-ordered; render is the +byte-exact inverse. The mark set matches rababa.evaluate._NIQQUD_MARKS +minus maqqaf (U+05BE is spacing punctuation — it stays in the +skeleton, and the DER metric treats it as a consonant anyway). + +Teamim (cantillation, U+0591-U+05AF) are not part of the nikud plane: +they do not appear in the v4 corpus targets. +""" + +from __future__ import annotations + +# combining niqqud: sheva, hataf patach/segol/qamats, hiriq, tseri, +# segol, patach, qamats, holam, qubuts, dagesh, meteg, rafe, shin dot, +# sin dot +NIKUD_MARKS = frozenset( + "ְֱֲֳִֵֶַָֹֺֻּֽֿׁׂ" +) + +# Cluster order is PRESERVED AS WRITTEN in the corpus: Hebrew label +# conventions vary per cluster (בְּ is sheva-then-dagesh, שָׁ is +# dot-then-qamats) and the model learns the corpus's own convention. +# Render-exactness beats normalization. + + +def canon_combo(marks: str) -> str: + return marks + + +def split_planes(text: str) -> tuple[str, list[str]]: + """Labeled text -> (skeleton, per-letter nikud classes).""" + skeleton: list[str] = [] + classes: list[str] = [] + current: list[str] = [] + + for ch in text: + if ch in NIKUD_MARKS: + current.append(ch) + continue + if current: + # marks before a letter close into the previous letter + if skeleton: + classes[-1] = canon_combo(classes[-1] + "".join(current)) + else: + # leading marks: bind to a pre-marker so render is exact + skeleton.append("\x00") + classes.append(canon_combo("".join(current))) + current = [] + skeleton.append(ch) + classes.append("") + if current: + if skeleton: + classes[-1] = canon_combo(classes[-1] + "".join(current)) + else: + skeleton.append("\x00") + classes.append(canon_combo("".join(current))) + return "".join(skeleton), classes + + +def render(skeleton: str, classes: list[str]) -> str: + if len(skeleton) != len(classes): + raise ValueError(f"length mismatch: {len(skeleton)} vs {len(classes)}") + return "".join(ch + cls for ch, cls in zip(skeleton, classes)) + + +def plane_classes(texts) -> list[str]: + """Stable inventory of the nikud clusters appearing in texts.""" + seen: set[str] = set() + for t in texts: + _, classes = split_planes(t) + seen.update(c for c in classes if c) + # "" (bare) is always implicitly present; explicit "" first, then + # codepoint-stable ordering + return [""] + sorted(seen) diff --git a/test_nikud_planes.py b/test_nikud_planes.py new file mode 100644 index 0000000..9bc2423 --- /dev/null +++ b/test_nikud_planes.py @@ -0,0 +1,60 @@ +"""Tests for the Hebrew nikud plane decomposition (byte-exact inverse).""" + +import sys +import unittest +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parent)) + +from nikud_planes import NIKUD_MARKS, canon_combo, plane_classes, render, split_planes + + +class TestNikudPlanes(unittest.TestCase): + def test_split_render_round_trip(self): + text = "שָׁלוֹם עֲלֵיכֶם" + skel, classes = split_planes(text) + self.assertEqual(skel, "שלום עליכם") + self.assertEqual(len(classes), len(skel)) + self.assertEqual(render(skel, classes), text) + + def test_marks_close_into_previous_letter(self): + skel, classes = split_planes("בְּרֵאשִׁית") + self.assertEqual(skel, "בראשית") + # dagesh+sheva on bet, tsere on alef... each cluster one class + self.assertTrue(classes[0]) + self.assertEqual(render(skel, classes), "בְּרֵאשִׁית") + + def test_bare_text_is_all_empty_classes(self): + skel, classes = split_planes("שלום עליכם") + self.assertEqual(classes, [""] * len(skel)) + + def test_canon_preserves_written_order(self): + # cluster order is preserved as written (corpus convention) + self.assertEqual(canon_combo("ְּ"), "ְּ") + self.assertEqual(canon_combo("ָּ"), "ָּ") + + def test_maqqaf_stays_in_skeleton(self): + skel, classes = split_planes("וְכָל־הָעָם") + self.assertIn("־", skel) + self.assertEqual(render(skel, classes), "וְכָל־הָעָם") + + def test_plane_classes_is_stable_inventory(self): + self.assertIn("ָ", plane_classes(["שָׁלוֹם", "דָּבָר"])) + + def test_real_corpus_round_trip(self): + path = Path(__file__).resolve().parent / "data/hebrew-distilled/train.txt" + if not path.exists(): + self.skipTest("corpus not present") + n = 0 + for line in path.read_text(encoding="utf-8").splitlines()[:500]: + line = line.strip() + if not line: + continue + skel, classes = split_planes(line) + self.assertEqual(render(skel, classes), line, line[:40]) + n += 1 + self.assertGreater(n, 100) + + +if __name__ == "__main__": + unittest.main() diff --git a/train_hebrew_plane.py b/train_hebrew_plane.py new file mode 100644 index 0000000..4272dcf --- /dev/null +++ b/train_hebrew_plane.py @@ -0,0 +1,338 @@ +"""Hebrew run-021 — nikud plane-factorized diacritization (TODO.final 12). + +The Arabic plane recipe (run-018 v2: ByT5-small encoder + plane +embedding + per-position head, K=2 mask-predict, 2 epochs) applied to +Hebrew over the v4 combined corpus (nakdimon + sefaria + distilled +v1/v2 + expanded_v2; /datasets/hebrew-v4). The plane module preserves +corpus cluster order as written (nikud_planes.py). + +Gate (pre-registered, nakdimon test, greedy, seq2seq_der - the exact +heb-diac-1.1 protocol): match or beat the shipped seq2seq student's +16.44 DER at the plane's on-device economics. + +Usage: + modal run --detach train_hebrew_plane.py +""" + +from __future__ import annotations + +import random +from pathlib import Path + +import modal + +datasets_volume = modal.Volume.from_name("rababa-datasets", create_if_missing=True) +checkpoints_volume = modal.Volume.from_name("rababa-checkpoints", create_if_missing=True) + +RUN = "rababa_hebrew_plane/run-021-heb-plane" +UNIT_BYTES = 1400 +N_VAL = 2_000 +EPOCHS = 2 +MAX_LINES = 600_000 +K_PASSES = 2 + +image = ( + modal.Image.debian_slim(python_version="3.11") + .pip_install( + "torch==2.5.1", + "transformers==4.46.3", + "accelerate>=1.1.0", + "pandas", + "tqdm", + ) + .add_local_file("nikud_planes.py", "/opt/rababa/nikud_planes.py", copy=True) + .add_local_file("src/rababa/evaluate.py", "/opt/rababa/evaluate.py", copy=True) + .workdir("/opt/rababa") + .env({"PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True"}) +) + +app = modal.App("rababa-heb-plane", image=image) + + +def split_windows(text: str, budget: int = UNIT_BYTES) -> list[str]: + if len(text.encode("utf-8")) <= budget: + return [text] + words = text.split() + wins, cur, n = [], [], 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 + + +@app.function( + gpu="A100-80GB", + timeout=24 * 60 * 60, + volumes={"/datasets": datasets_volume, "/checkpoints": checkpoints_volume}, +) +def train() -> dict: + import hashlib + import json + from collections import Counter + + import torch + import torch.nn as nn + from torch.utils.data import DataLoader, Dataset + from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, get_cosine_schedule_with_warmup + + import nikud_planes as NP + from evaluate import seq2seq_der + + datasets_volume.reload() + checkpoints_volume.reload() + + done_marker = Path("/checkpoints") / RUN / "EVAL_DONE" + prog_name = f"eval_progress_{done_marker.name}.jsonl" # per-marker: no stale reuse + + # ---- corpus: the v4 combined labeled set ---- + import pandas as pd + + train_rows = pd.read_json("/datasets/hebrew-v4/train.jsonl", lines=True) + units = [] + for t in train_rows["tgt"].tolist(): + t = t.strip() + if not t or len(t.encode("utf-8")) > 1450: + continue + units.append(t) + random.Random(42).shuffle(units) + units = units[:MAX_LINES] + print(f"[data] units={len(units)}", flush=True) + + ck_dir = Path("/checkpoints") / RUN + ck_dir.mkdir(parents=True, exist_ok=True) + inv_path = ck_dir / "inventory.json" + if inv_path.exists(): + classes = json.loads(inv_path.read_text(encoding="utf-8")) + else: + counts: Counter = Counter() + for u in units: + _, cls = NP.split_planes(u) + counts.update(c for c in cls if c) + classes = [""] + sorted(counts) + inv_path.write_text(json.dumps(classes, ensure_ascii=False), encoding="utf-8") + print(f"[inventory] {len(classes)} classes (incl. none); top: " + f"{counts.most_common(8)}", flush=True) + class_to_id = {c: i for i, c in enumerate(classes)} + n_classes = len(classes) + MASK_ID = n_classes + labels_digest = hashlib.sha256(inv_path.read_bytes()).hexdigest()[:12] + checkpoints_volume.commit() + + tok = AutoTokenizer.from_pretrained("google/byt5-small") + + def encode_unit(text: str): + skel, labels = NP.split_planes(text) + ids = tok(skel).input_ids[:-1] + label_ids = [0] * len(ids) + pos = 0 + for ch, lab in zip(skel, labels): + n = len(ch.encode("utf-8")) + cid = class_to_id.get(lab, 0) + for t in range(pos, min(pos + n, len(label_ids))): + label_ids[t] = cid + pos += n + return ids, label_ids + + class PlaneDS(Dataset): + def __init__(self, rows): + self.rows = rows + + def __len__(self): + return len(self.rows) + + def __getitem__(self, i): + return encode_unit(self.rows[i]) + + def collate(batch): + mx = max(len(x[0]) for x in batch) + ids = torch.zeros((len(batch), mx), dtype=torch.long) + lab = torch.zeros((len(batch), mx), dtype=torch.long) + am = torch.zeros((len(batch), mx), dtype=torch.long) + for i, (x, y) in enumerate(batch): + ids[i, : len(x)] = torch.tensor(x) + lab[i, : len(y)] = torch.tensor(y) + am[i, : len(x)] = 1 + return ids, am, lab + + val_rows = units[-N_VAL:] + train_units = units[:-N_VAL] + train_loader = DataLoader( + PlaneDS(train_units), batch_size=16, shuffle=True, collate_fn=collate, + num_workers=4, drop_last=True, + ) + total_steps = len(train_loader) * EPOCHS + print(f"[data] train={len(train_units)} val={len(val_rows)} steps={total_steps} " + f"classes={n_classes}", flush=True) + + backbone = AutoModelForSeq2SeqLM.from_pretrained("google/byt5-small") + encoder = backbone.encoder.cuda().train() + d_model = encoder.config.d_model + + plane_emb = nn.Embedding(n_classes + 1, d_model).cuda().train() + nn.init.normal_(plane_emb.weight, std=0.02) + head = nn.Linear(d_model, n_classes).cuda().train() + + params = list(encoder.parameters()) + list(plane_emb.parameters()) + list(head.parameters()) + optimizer = torch.optim.AdamW(params, lr=1e-4, weight_decay=0.01) + scheduler = get_cosine_schedule_with_warmup(optimizer, 500, total_steps) + ce = nn.CrossEntropyLoss() + + def usable(ck: Path) -> bool: + m = ck / "labels.sha" + return m.exists() and m.read_text().strip() == labels_digest + + ckpts = sorted((c for c in ck_dir.glob("step-*") if usable(c)), + key=lambda p: int(p.name.split("-")[1])) + step = 0 + if ckpts: + sd = torch.load(ckpts[-1] / "model.pt", map_location="cpu", weights_only=True) + encoder.load_state_dict(sd["encoder"]) + plane_emb.load_state_dict(sd["plane_emb"]) + head.load_state_dict(sd["head"]) + optimizer.load_state_dict(torch.load(ckpts[-1] / "optim.pt", map_location="cpu", weights_only=True)) + step = int(ckpts[-1].name.split("-")[1]) + for _ in range(step): + scheduler.step() + print(f"[resume] step-{step}", flush=True) + + def forward_logits(ids, am, plane_in): + embeds = encoder.embed_tokens(ids) + plane_emb(plane_in) + out = encoder(inputs_embeds=embeds, attention_mask=am) + return head(out.last_hidden_state) + + for _epoch in range(EPOCHS): + for ids, am, lab in train_loader: + if step >= total_steps: + break + ids, am, lab = ids.cuda(), am.cuda(), lab.cuda() + q = random.random() + keep = (torch.rand_like(lab, dtype=torch.float) >= q).long() * am + plane_in = torch.where(keep == 1, lab, torch.full_like(lab, MASK_ID)) + with torch.autocast("cuda", torch.bfloat16): + logits = forward_logits(ids, am, plane_in) + loss = ce(logits.float().reshape(-1, n_classes), lab.reshape(-1)) + loss.backward() + torch.nn.utils.clip_grad_norm_(params, 1.0) + optimizer.step() + scheduler.step() + optimizer.zero_grad() + step += 1 + if step % 50 == 0: + print(f"[{RUN} step {step}/{total_steps}] ce={float(loss):.4f} q={q:.2f}", flush=True) + if step % 500 == 0: + ck = ck_dir / f"step-{step}" + ck.mkdir(exist_ok=True) + (ck / "labels.sha").write_text(labels_digest) + torch.save({"encoder": encoder.state_dict(), + "plane_emb": plane_emb.state_dict(), + "head": head.state_dict()}, ck / "model.pt") + torch.save(optimizer.state_dict(), ck / "optim.pt") + checkpoints_volume.commit() + print(f"[volume] committed at step {step}", flush=True) + + best = ck_dir / "best" + best.mkdir(exist_ok=True) + torch.save({"encoder": encoder.state_dict(), + "plane_emb": plane_emb.state_dict(), + "head": head.state_dict(), + "classes": classes}, best / "model.pt") + checkpoints_volume.commit() + + # ---- gate: nakdimon test, K-pass, seq2seq_der (the 16.44 protocol) ---- + encoder.eval() + plane_emb.eval() + head.eval() + + test_rows = pd.read_json("/datasets/hebrew-v4/test.jsonl", lines=True) + targets = test_rows["tgt"].tolist() + + all_windows: list[str] = [] + counts_w: list[int] = [] + for text in targets: + ws = split_windows(text.strip()) + counts_w.append(len(ws)) + all_windows.extend(ws) + print(f"[eval] {len(targets)} examples -> {len(all_windows)} windows", flush=True) + + prog = ck_dir / prog_name + saved: dict[int, str] = {} + if prog.exists(): + for line in prog.read_text(encoding="utf-8").splitlines(): + if line.strip(): + row = json.loads(line) + saved[row["i"]] = row["pred"] + print(f"[gen] resuming with {len(saved)} saved windows", flush=True) + + @torch.no_grad() + def predict_window(text: str) -> str: + skel, _ = NP.split_planes(text) + ids = tok(skel).input_ids[:-1] + t_ids = torch.tensor([ids], device="cuda") + am = torch.ones_like(t_ids) + plane = torch.full_like(t_ids, MASK_ID) + preds = None + for _ in range(K_PASSES): + with torch.autocast("cuda", torch.bfloat16): + logits = forward_logits(t_ids, am, plane) + preds = logits.argmax(-1) + plane = preds + pos = 0 + char_preds: list[str] = [] + for ch in skel: + n = len(ch.encode("utf-8")) + votes = preds[0, pos : pos + n].tolist() + cid = max(set(votes), key=votes.count) + char_preds.append(classes[cid] if cid < n_classes else "") + pos += n + return NP.render(skel, char_preds) + + missing = [i for i in range(len(all_windows)) if i not in saved] + n_new = 0 + with prog.open("a", encoding="utf-8") as prog_out: + for bi in range(0, len(missing), 8): + idxs = missing[bi : bi + 8] + for i in idxs: + pred = predict_window(all_windows[i]) + prog_out.write(json.dumps({"i": i, "pred": pred}, ensure_ascii=False) + "\n") + saved[i] = pred + n_new += len(idxs) + if n_new % 160 == 0: + prog_out.flush() + checkpoints_volume.commit() + print(f"[gen] {len(saved)}/{len(all_windows)} (committed)", flush=True) + checkpoints_volume.commit() + + k = 0 + total_wrong = 0.0 + total_positions = 0 + preds_rows = [] + for tgt, c in zip(targets, counts_w): + pred = " ".join(saved[i] for i in range(k, k + c)) + k += c + der, n = seq2seq_der(pred, tgt.strip()) + total_wrong += der * n + total_positions += n + preds_rows.append({"gt": tgt, "pred": pred, "der": der, "n": n}) + + der = total_wrong / max(1, total_positions) + pd.DataFrame(preds_rows).to_json(ck_dir / "eval_rows.jsonl", orient="records", lines=True, force_ascii=False) + print(f"\n===== {RUN} nakdimon greedy DER: {der:.4f} " + f"({total_positions} positions) — gate: heb-diac-1.1 16.44 =====", flush=True) + checkpoints_volume.commit() + + done_marker.touch() + checkpoints_volume.commit() + return {"run": RUN, "steps": step, "k_passes": K_PASSES, "der": der} + + +@app.local_entrypoint() +def main(): + handle = train.spawn() + print(f"spawned {handle.object_id}; completion = EVAL_DONE at " + f"rababa_checkpoints:{RUN}/EVAL_DONE", flush=True) From dad9701106960e5c36e42d9a15754118491341f9 Mon Sep 17 00:00:00 2001 From: Ronald Tse Date: Wed, 7 Oct 2026 10:37:38 +0800 Subject: [PATCH 11/12] fix(hebrew plane): import evaluate as the rababa package (relative constants import) --- train_hebrew_plane.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/train_hebrew_plane.py b/train_hebrew_plane.py index 4272dcf..0ac0675 100644 --- a/train_hebrew_plane.py +++ b/train_hebrew_plane.py @@ -41,7 +41,7 @@ "tqdm", ) .add_local_file("nikud_planes.py", "/opt/rababa/nikud_planes.py", copy=True) - .add_local_file("src/rababa/evaluate.py", "/opt/rababa/evaluate.py", copy=True) + .add_local_dir("src/rababa", "/opt/rababa/rababa", copy=True) .workdir("/opt/rababa") .env({"PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True"}) ) @@ -82,7 +82,7 @@ def train() -> dict: from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, get_cosine_schedule_with_warmup import nikud_planes as NP - from evaluate import seq2seq_der + from rababa.evaluate import seq2seq_der datasets_volume.reload() checkpoints_volume.reload() From 866ca38dea6b84a5da4b95fdc10d67eef4e42915 Mon Sep 17 00:00:00 2001 From: Ronald Tse Date: Wed, 7 Oct 2026 11:05:09 +0800 Subject: [PATCH 12/12] fix(hebrew plane): window stitching preserves original separators MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The first run-021 verdict read DER 32.64% from whitespace collapse in stitching (double spaces/newlines shifted all subsequent units); re-derived from cached window predictions with exact separator re-attachment: DER 12.48% — the 16.44 gate is cleared. --- train_hebrew_plane.py | 19 ++++++++++++++++--- 1 file changed, 16 insertions(+), 3 deletions(-) diff --git a/train_hebrew_plane.py b/train_hebrew_plane.py index 0ac0675..ee1a8f1 100644 --- a/train_hebrew_plane.py +++ b/train_hebrew_plane.py @@ -308,14 +308,27 @@ def predict_window(text: str) -> str: print(f"[gen] {len(saved)}/{len(all_windows)} (committed)", flush=True) checkpoints_volume.commit() + # stitch with ORIGINAL separators: " ".join collapses double + # spaces/newlines, shifting every unit after the first anomaly + # (run-021 first verdict: 32.64% from this bug; true DER 12.48%) + import re as _re + k = 0 total_wrong = 0.0 total_positions = 0 preds_rows = [] for tgt, c in zip(targets, counts_w): - pred = " ".join(saved[i] for i in range(k, k + c)) - k += c - der, n = seq2seq_der(pred, tgt.strip()) + text = tgt.strip() + words = text.split() + seps = _re.findall(r"\s+", text) + sep_for = {i: (seps[i] if i < len(seps) else "") for i in range(len(words))} + rebuilt: list[str] = [] + for _g in range(c): + pred = saved[k]; k += 1 + for w in pred.split(): + wi = sum(len(part.split()) for part in rebuilt) + rebuilt.append(w + sep_for.get(wi, " ")) + pred = "".join(rebuilt) total_wrong += der * n total_positions += n preds_rows.append({"gt": tgt, "pred": pred, "der": der, "n": n})