diff --git a/export_plane_onnx_hf.py b/export_plane_onnx_hf.py new file mode 100644 index 0000000..b35b9f2 --- /dev/null +++ b/export_plane_onnx_hf.py @@ -0,0 +1,128 @@ +# /// script +# requires-python = ">=3.11" +# dependencies = [ +# "torch==2.5.1", +# "transformers==4.46.3", +# "onnx==1.17.0", +# "onnxruntime==1.20.1", +# "numpy<2", +# ] +# /// +"""WO04: export run-022 best checkpoint to plane.onnx (+dynamic QInt8). + +Reads /ckpt/run-022-heb-plane-base/best/model.pt, exports fp32, +quantizes (same recipe as the ara plane family), sanity-decodes a few +windows on CPU ORT, writes artifacts + classes.json to the bucket. + +Usage: + hf jobs uv run --flavor l4x1 -d --timeout 2h \ + -v hf://buckets/Interscript/isx-training:/ckpt:rw \ + -v /tmp/isx-hf-code:/code:ro export_plane_onnx_hf.py +""" + +from __future__ import annotations + +import json +import sys +from pathlib import Path + +sys.path.insert(0, "/code") +sys.path.insert(0, "/code/src") + +CKPT = Path("/ckpt/run-022-heb-plane-base") +OUT = CKPT / "export" +K_PASSES = 3 + +import torch +import torch.nn as nn +from transformers import AutoTokenizer + + +def main() -> None: + OUT.mkdir(parents=True, exist_ok=True) + sd = torch.load(CKPT / "best/model.pt", map_location="cpu", weights_only=True) + classes = sd["classes"] + n_classes = len(classes) + + from transformers import AutoModelForSeq2SeqLM + + backbone = AutoModelForSeq2SeqLM.from_pretrained("google/byt5-base") + encoder = backbone.encoder.eval() + d_model = encoder.config.d_model + plane_emb = nn.Embedding(n_classes + 1, d_model) + head = nn.Linear(d_model, n_classes) + encoder.load_state_dict(sd["encoder"]) + plane_emb.load_state_dict(sd["plane_emb"]) + head.load_state_dict(sd["head"]) + + class PlaneGraph(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) + return self.head(out.last_hidden_state) # class_logits + + graph = PlaneGraph().eval() + ids = torch.zeros((1, 64), dtype=torch.long) + plane = torch.full_like(ids, n_classes) + fp32_path = OUT / "plane-base-fp32.onnx" + with torch.no_grad(): + torch.onnx.export( + graph, (ids, plane), str(fp32_path), + 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=17, dynamo=False, + ) + print(f"[export] fp32 {fp32_path.stat().st_size / 1e6:.1f} MB", flush=True) + + from onnxruntime.quantization import QuantType, quantize_dynamic + + int8_path = OUT / "plane-base-int8.onnx" + quantize_dynamic(str(fp32_path), str(int8_path), weight_type=QuantType.QInt8) + print(f"[export] int8 {int8_path.stat().st_size / 1e6:.1f} MB", flush=True) + + # CPU sanity: decode a couple of real test windows with the runtime path + import numpy as np + import onnxruntime as ort + + import nikud_planes as NP + + sess = ort.InferenceSession(str(int8_path), providers=["CPUExecutionProvider"]) + tok = AutoTokenizer.from_pretrained("google/byt5-base") + test = [json.loads(l) for l in Path("/code/../dev/null").read_text()] if False else None + samples = [ + "בראשית ברא אלהים את השמים ואת הארץ", + "וזה לא יספיק לבית הספר בעיר", + ] + for text in samples: + skel, _ = NP.split_planes(text) + ids = tok(skel).input_ids[:-1] + t_ids = np.array([ids], dtype=np.int64) + plane = np.full_like(t_ids, n_classes) + for _ in range(K_PASSES): + logits = sess.run(None, {"input_ids": t_ids, "plane_ids": plane})[0] + plane = logits.argmax(-1) + pos, char_preds = 0, [] + for ch in skel: + n = len(ch.encode("utf-8")) + votes = plane[0, pos:pos + n].tolist() + cid = max(set(votes), key=votes.count) + char_preds.append(classes[cid] if cid < n_classes else "") + pos += n + print(f"[sanity] {text[:20]} -> {NP.render(skel, char_preds)[:40]}", flush=True) + + (OUT / "classes.json").write_text(json.dumps(classes, ensure_ascii=False), encoding="utf-8") + (OUT / "EXPORT_DONE").write_text("ok\n") + print("[done] EXPORT_DONE", flush=True) + + +if __name__ == "__main__": + main() diff --git a/label_arwiki_hf.py b/label_arwiki_hf.py new file mode 100644 index 0000000..d736baf --- /dev/null +++ b/label_arwiki_hf.py @@ -0,0 +1,125 @@ +# /// script +# requires-python = ">=3.11" +# dependencies = [ +# "torch==2.5.1", +# "transformers==4.46.3", +# "accelerate>=1.1.0", +# "pyarrow", +# ] +# /// +"""WO06 stage 1: r7 teacher pseudo-labels arwiki (noisy student). + +NO LLM LABELS — the teacher is our own r7 dedicated model (run-007-news). +Greedy decode (the protocol decode) + a teacher-forced margin pass; +windows are kept when a high fraction of decoded tokens are confident +(top1-top2 logit margin > 1.0 nat). Output: arwiki-pseudo.jsonl. + +Usage: + hf jobs uv run --flavor a10g-large -d --timeout 8h \ + -v hf://datasets/Interscript/arabic-r8:/train_data:ro \ + -v hf://buckets/Interscript/isx-training:/ckpt:rw \ + label_arwiki_hf.py +""" + +from __future__ import annotations + +import json +import time +from pathlib import Path + +import torch +from transformers import AutoModelForSeq2SeqLM, AutoTokenizer + +DATA = Path("/train_data") +CKPT = Path("/ckpt/r8-pseudo") +TEACHER = Path("/ckpt/r7-best") +UNIT_BYTES = 1400 +MARGIN_NAT = 1.0 +KEEP_FRAC = 0.9 +MAX_UNITS = 400_000 + + +def window(text: str, budget: int = UNIT_BYTES) -> list[str]: + if len(text.encode("utf-8")) <= budget: + return [text] + words, cur, n, wins = text.split(), [], 0, [] + for w in words: + c = len(w.encode("utf-8")) + 1 + if cur and n + c > budget: + wins.append(" ".join(cur)); cur, n = [], 0 + cur.append(w); n += c + if cur: + wins.append(" ".join(cur)) + return wins + + +def main() -> None: + CKPT.mkdir(parents=True, exist_ok=True) + device = "cuda" + tok = AutoTokenizer.from_pretrained(TEACHER) + model = AutoModelForSeq2SeqLM.from_pretrained( + TEACHER, torch_dtype=torch.bfloat16).to(device).eval() + + units: list[str] = [] + for name in ("train.txt", "val.txt"): + p = DATA / name + if not p.exists(): + continue + for line in p.read_text(encoding="utf-8").splitlines(): + line = line.strip() + if not line: + continue + units.extend(window(line)) + if len(units) >= MAX_UNITS: + break + if len(units) >= MAX_UNITS: + break + print(f"[data] arwiki units={len(units)}", flush=True) + + out_path = CKPT / "arwiki-pseudo.jsonl" + done = 0 + kept = 0 + t0 = time.time() + with out_path.open("w", encoding="utf-8") as out, torch.no_grad(): + for i in range(0, len(units), 8): + batch = units[i:i + 8] + enc = tok(batch, return_tensors="pt", padding=True, truncation=True, + max_length=1450).to(device) + with torch.autocast("cuda", torch.bfloat16): + gen = model.generate(**enc, max_new_tokens=1500, num_beams=1, + return_dict_in_generate=True, output_scores=False) + seqs = gen.sequences + texts = tok.batch_decode(seqs, skip_special_tokens=True) + + # teacher-forced margin pass on the generated tokens + dec = seqs[:, :-1] + dec_mask = dec != tok.pad_token_id + with torch.autocast("cuda", torch.bfloat16): + sc = model(**{k: enc[k] for k in ("input_ids", "attention_mask")}, + decoder_input_ids=dec, decoder_attention_mask=dec_mask) + logits = sc.logits.float() + top2 = logits.topk(2, dim=-1).values + margins = (top2[..., 0] - top2[..., 1]) + conf = (margins > MARGIN_NAT) & dec_mask + frac = conf.sum(-1).float() / dec_mask.sum(-1).clamp(min=1) + + srcs = tok.batch_decode(enc.input_ids, skip_special_tokens=True) + for src, text, f in zip(srcs, texts, frac.tolist()): + done += 1 + if f >= KEEP_FRAC: + kept += 1 + out.write(json.dumps({"src": src, "tgt": text, + "keep_frac": round(f, 4)}, + ensure_ascii=False) + "\n") + if (i // 8) % 50 == 0: + out.flush() + rate = done / max(1, time.time() - t0) + print(f"[gen] {done}/{len(units)} kept={kept} " + f"({rate:.1f} win/s)", flush=True) + + print(f"[done] kept {kept}/{done} windows -> {out_path}", flush=True) + (CKPT / "LABEL_DONE").write_text(f"kept={kept} total={done}\n") + + +if __name__ == "__main__": + main() diff --git a/probe_hebrew_audio_pseudolabel.py b/probe_hebrew_audio_pseudolabel.py new file mode 100644 index 0000000..b0e3892 --- /dev/null +++ b/probe_hebrew_audio_pseudolabel.py @@ -0,0 +1,92 @@ +# /// script +# requires-python = ">=3.11" +# dependencies = [ +# "torch==2.5.1", +# "transformers==4.46.3", +# "librosa>=0.10", +# "numpy<2", +# ] +# /// +"""WO09 (PARKED probe): Hebrew audio -> phoneme pseudo-label viability. + +Stage 1: universal phoneme CTC ASR (facebook/wav2vec2-lv-60-espeak-cv-ft) +over a small local dir of Hebrew audio (the script fetches NOTHING). +Stage 2: forced-alignment viability — espeak-style reference phonemes of +our plane-restored nikud text vs the ASR phoneme stream; token agreement +stats + per-hour labeling cost estimate. Output: a go/no-go report for a +2027 audio-supervision campaign (the ReNikud direction, arXiv 2606.20179). + +Usage: + hf jobs uv run --flavor l4x1 --timeout 1h \ + -v hf://buckets/Interscript/isx-training:/ckpt:rw \ + probe_hebrew_audio_pseudolabel.py --audio-dir /ckpt/audio-sample +""" + +from __future__ import annotations + +import argparse +import json +import time +from pathlib import Path + +ASR = "facebook/wav2vec2-lv-60-espeak-cv-ft" + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--audio-dir", required=True) + ap.add_argument("--limit", type=int, default=20) + ap.add_argument("--out", default="probe_report.json") + args = ap.parse_args() + + audio_dir = Path(args.audio_dir) + clips = sorted(p for p in audio_dir.iterdir() + if p.suffix.lower() in (".wav", ".flac", ".mp3"))[:args.limit] + if not clips: + print(f"no audio clips under {audio_dir} — probe requires local audio " + f"(parked WO; nothing is fetched).", flush=True) + return + + import librosa + import torch + from transformers import AutoProcessor, AutoModelForCTC + + processor = AutoProcessor.from_pretrained(ASR) + model = AutoModelForCTC.from_pretrained(ASR).eval() + if torch.cuda.is_available(): + model = model.cuda() + + results = [] + t0 = time.time() + for clip in clips: + speech, _ = librosa.load(str(clip), sr=16000, mono=True) + inputs = processor(speech, sampling_rate=16000, return_tensors="pt") + if torch.cuda.is_available(): + inputs = {k: v.cuda() for k, v in inputs.items()} + with torch.no_grad(): + logits = model(**inputs).logits + ids = logits.argmax(-1) + phon = processor.batch_decode(ids) + results.append({"clip": clip.name, "phonemes": phon[0], + "seconds": round(len(speech) / 16000, 2)}) + + wall = time.time() - t0 + total_s = sum(r["seconds"] for r in results) or 1.0 + report = { + "status": "POC: ASR stage runs; alignment stage NOT implemented (parked)", + "clips": len(results), + "audio_seconds": total_s, + "wall_s": round(wall, 1), + "rtf_asr_only": round(wall / total_s, 3), + "hourly_label_cost_usd": round(2.50 * wall / total_s, 2), + "next_steps": "espeak reference phonemization of plane-restored text; " + "CTC forced alignment; agreement stats; 2027 campaign gate", + "samples": results[:3], + } + Path(args.out).write_text(json.dumps(report, ensure_ascii=False, indent=1)) + print(json.dumps({k: v for k, v in report.items() if k != "samples"}, + ensure_ascii=False, indent=1), flush=True) + + +if __name__ == "__main__": + main() diff --git a/scripts/build_hebrew_v4_local.py b/scripts/build_hebrew_v4_local.py new file mode 100644 index 0000000..59c7a10 --- /dev/null +++ b/scripts/build_hebrew_v4_local.py @@ -0,0 +1,87 @@ +"""Rebuild the hebrew-v4 combined corpus locally (no Modal). + +Mirrors train_hebrew_v4.py::build_combined_corpus exactly — same +sources (local data/ dirs), same strip/dedupe/filter order — so the +HF-mounted dataset is byte-comparable to the run-021 volume corpus. +""" + +from __future__ import annotations + +import json +import sys +from pathlib import Path + +ROOT = Path(__file__).resolve().parent.parent +DATA = ROOT / "data" + +_NIKUD_MARKS = set("ְֱֲֳִֵֶַָֹֺֻּֽֿׁׂ־") + + +def _strip_nikud(s: str) -> str: + return "".join(c for c in s if c not in _NIKUD_MARKS) + + +def _has_consonants(s: str) -> bool: + hebrew_consonants = set("אבגדהוזחטיכלמנסעפצקרשתךםןףץ") + return any(c in hebrew_consonants for c in s) + + +def _load_labeled_lines(path: Path) -> list[str]: + if not path.is_file(): + return [] + out = [] + for line in path.read_text(encoding="utf-8").splitlines(): + line = line.strip() + if line and any("֑" <= c <= "ׇ" for c in line): + out.append(line) + return out + + +def build(out_root: Path) -> dict: + sources = { + "nakdimon": DATA / "nakdimon/train.txt", + "sefaria_tanakh": DATA / "sefaria-tanakh/train.txt", + "distilled_v1": DATA / "hebrew-distilled/train.txt", + "distilled_v2": DATA / "hebrew-dictabert-distilled/train.txt", + "expanded_v2": DATA / "hebrew-expanded-v2/train.txt", + } + all_pairs = [] + seen: set[tuple[str, str]] = set() + counts = {} + for name, path in sources.items(): + new_count = 0 + for line in _load_labeled_lines(path): + undiacritized = _strip_nikud(line).strip() + if not undiacritized or not _has_consonants(undiacritized): + continue + if len(undiacritized) < 5 or len(undiacritized) > 500: + continue + key = (undiacritized, line) + if key in seen: + continue + seen.add(key) + all_pairs.append({"src": undiacritized, "tgt": line}) + new_count += 1 + counts[name] = new_count + + out_root.mkdir(parents=True, exist_ok=True) + with (out_root / "train.jsonl").open("w", encoding="utf-8") as f: + for pair in all_pairs: + f.write(json.dumps(pair, ensure_ascii=False) + "\n") + + test_pairs = [] + for line in _load_labeled_lines(DATA / "nakdimon/test.txt"): + undiacritized = _strip_nikud(line).strip() + if undiacritized and _has_consonants(undiacritized): + test_pairs.append({"src": undiacritized, "tgt": line}) + with (out_root / "test.jsonl").open("w", encoding="utf-8") as f: + for pair in test_pairs: + f.write(json.dumps(pair, ensure_ascii=False) + "\n") + + return {"train": len(all_pairs), "test": len(test_pairs), "counts": counts} + + +if __name__ == "__main__": + out = Path(sys.argv[1]) if len(sys.argv) > 1 else DATA / "hebrew-v4" + stats = build(out) + print(json.dumps(stats, ensure_ascii=False, indent=1)) diff --git a/train_arabic_r8_hf.py b/train_arabic_r8_hf.py new file mode 100644 index 0000000..d390a69 --- /dev/null +++ b/train_arabic_r8_hf.py @@ -0,0 +1,274 @@ +# /// script +# requires-python = ">=3.11" +# dependencies = [ +# "torch==2.5.1", +# "transformers==4.46.3", +# "accelerate>=1.1.0", +# "pyarrow", +# "pandas", +# "pyarabic", +# "prettytable", +# ] +# /// +"""WO06 stage 2: r8 noisy-student — gold tashkeela + r7-pseudo arwiki. + +Student initializes FROM r7 (born-again variant: guarantees the r7 +starting point) and trains one epoch on the interleaved mixture. +Gate (pre-registered): beat 2.2864 DER on SadeedDiac-25 under the exact +windowed protocol (WINDOW=600, greedy, project_haraqat, Misraj evaluator, +gt_missing_diacritic_is_error=False). + +Usage: + hf jobs uv run --flavor a100-large -d --timeout 20h \ + -v hf://datasets/Interscript/arabic-r8:/train_data:ro \ + -v hf://buckets/Interscript/isx-training:/ckpt:rw \ + -v /tmp/isx-hf-code:/code:ro train_arabic_r8_hf.py +""" + +from __future__ import annotations + +import json +import random +import re +import sys +import time +from pathlib import Path + +sys.path.insert(0, "/code") + +DATA = Path("/train_data") +CKPT = Path("/ckpt/run-024-arabic-r8-noisy") +TEACHER = Path("/ckpt/r7-best") +RUN = "run-024-arabic-r8-noisy" +WINDOW = 600 +GOLD_CAP = 400_000 +PSEUDO_CAP = 200_000 +MICRO_BS = 16 +GRAD_ACCUM = 2 +EPOCHS = 1 +LR = 5e-5 +WARMUP = 500 + +DIACRITICS_RE = re.compile("[ؐ-ًؚ-ٰٟۖ-ۜ۟-۪ۨ-ۭ]") + + +def split_windows(text: str, budget: int = WINDOW) -> list[str]: + if len(text.encode("utf-8")) <= budget: + return [text] + words, cur, n, wins = text.split(), [], 0, [] + for w in words: + c = len(w.encode("utf-8")) + 1 + if cur and n + c > budget: + wins.append(" ".join(cur)); cur, n = [], 0 + cur.append(w); n += c + if cur: + wins.append(" ".join(cur)) + return wins + + +def project_haraqat(pred: str, text: str) -> str: + from difflib import SequenceMatcher + + pred_haraqat = [""] + for ch in pred: + if DIACRITICS_RE.match(ch): + pred_haraqat[-1] += ch + else: + pred_haraqat.append("") + pred_haraqat = pred_haraqat[1:] + pred_letters = [c for c in pred if not DIACRITICS_RE.match(c)] + text_letters = [c for c in text if not DIACRITICS_RE.match(c)] + sm = SequenceMatcher(None, text_letters, pred_letters, autojunk=False) + out = [] + for op, i1, i2, j1, j2 in sm.get_opcodes(): + if op == "equal": + for k in range(i2 - i1): + out.append(text_letters[i1 + k] + pred_haraqat[j1 + k]) + else: + for k in range(i1, i2): + out.append(text_letters[k]) + return "".join(out) + + +def main() -> None: + import torch + from torch.utils.data import DataLoader, Dataset + from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, get_cosine_schedule_with_warmup + + CKPT.mkdir(parents=True, exist_ok=True) + rng = random.Random(42) + + gold: list[tuple[str, str]] = [] + for name in ("tashkeela-full/train-001.txt", "tashkeela-full/train-002.txt", + "tashkeela-full/train-003.txt", "tashkeela-full/val-001.txt"): + p = DATA / name + if not p.exists(): + continue + for line in p.read_text(encoding="utf-8").splitlines(): + tgt = line.strip() + if not tgt: + continue + src = DIACRITICS_RE.sub("", tgt).strip() + if len(src) < 5 or len(src.encode()) > WINDOW - 50: + continue + gold.append((src, tgt)) + if len(gold) >= GOLD_CAP: + break + if len(gold) >= GOLD_CAP: + break + print(f"[data] gold={len(gold)}", flush=True) + + pseudo: list[tuple[str, str]] = [] + ppath = Path("/ckpt/r8-pseudo/arwiki-pseudo.jsonl") + if ppath.exists(): + for line in ppath.read_text(encoding="utf-8").splitlines(): + if not line.strip(): + continue + r = json.loads(line) + src = DIACRITICS_RE.sub("", r["tgt"]).strip() + if len(src) < 5 or len(src.encode()) > WINDOW - 50: + continue + pseudo.append((src, r["tgt"])) + if len(pseudo) >= PSEUDO_CAP: + break + print(f"[data] pseudo={len(pseudo)}", flush=True) + + units = gold + pseudo + rng.shuffle(units) + rng.shuffle(gold) + + tok = AutoTokenizer.from_pretrained(TEACHER) + + class DS(Dataset): + def __getitem__(self, i): + src, tgt = units[i] + e = tok(src, truncation=True, max_length=WINDOW) + d = tok(tgt, truncation=True, max_length=WINDOW * 2) + return {"input_ids": e.input_ids, "attention_mask": e.attention_mask, + "labels": d.input_ids} + + def collate(b): + mx = max(len(x["input_ids"]) for x in b) + md = max(len(x["labels"]) for x in b) + ids = torch.zeros((len(b), mx), dtype=torch.long) + am = torch.zeros((len(b), mx), dtype=torch.long) + lb = torch.full((len(b), md), -100, dtype=torch.long) + for i, x in enumerate(b): + ids[i, :len(x["input_ids"])] = torch.tensor(x["input_ids"]) + am[i, :len(x["attention_mask"])] = torch.tensor(x["attention_mask"]) + lb[i, :len(x["labels"])] = torch.tensor(x["labels"]) + return ids, am, lb + + loader = DataLoader(DS(), batch_size=MICRO_BS, shuffle=True, collate_fn=collate, + num_workers=2, drop_last=True) + total_steps = (len(loader) // GRAD_ACCUM) * EPOCHS + print(f"[data] units={len(units)} steps={total_steps}", flush=True) + + model = AutoModelForSeq2SeqLM.from_pretrained(TEACHER, torch_dtype=torch.float32) + model = model.cuda().train() + model.gradient_checkpointing_enable() + model.config.use_cache = False + + optimizer = torch.optim.AdamW(model.parameters(), lr=LR, weight_decay=0.01) + scheduler = get_cosine_schedule_with_warmup(optimizer, WARMUP, total_steps) + + # resume contract + ckpts = sorted(CKPT.glob("step-*"), 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) + model.load_state_dict(sd) + 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 save(s: int) -> None: + d = CKPT / f"step-{s}" + d.mkdir(exist_ok=True) + torch.save(model.state_dict(), d / "model.pt") + torch.save(optimizer.state_dict(), d / "optim.pt") + print(f"[ckpt] step-{s}", flush=True) + + optimizer.zero_grad() + micro = 0 + t0 = time.time() + for _epoch in range(EPOCHS): + for ids, am, lb in loader: + if step >= total_steps: + break + ids, am, lb = ids.cuda(), am.cuda(), lb.cuda() + with torch.autocast("cuda", torch.bfloat16): + out = model(input_ids=ids, attention_mask=am, labels=lb) + (out.loss / GRAD_ACCUM).backward() + micro += 1 + if micro % GRAD_ACCUM == 0: + torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) + optimizer.step() + scheduler.step() + optimizer.zero_grad() + step += 1 + if step % 100 == 0: + print(f"[{RUN} {step}/{total_steps}] loss={float(out.loss):.4f} " + f"({step / (time.time() - t0):.2f} it/s)", flush=True) + if step % 1000 == 0: + save(step) + + best = CKPT / "best" + best.mkdir(exist_ok=True) + model.config.use_cache = True + model.save_pretrained(best) + tok.save_pretrained(best) + print("[ckpt] best saved", flush=True) + + # ---- gate: SadeedDiac-25, exact windowed protocol ---- + import pandas as pd + import pyarrow.parquet as pq + + model.eval() + table = pq.read_table(DATA / "sadeeddiac-25.parquet") + inputs = [DIACRITICS_RE.sub("", t) for t in table.column("input").to_pylist()] + outputs = table.column("output").to_pylist() + all_windows, counts = [], [] + 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) + + preds = [] + with torch.no_grad(): + for i in range(0, len(all_windows), 32): + batch = all_windows[i:i + 32] + enc = tok(batch, return_tensors="pt", padding=True, truncation=True, + max_length=WINDOW).to("cuda") + with torch.autocast("cuda", torch.bfloat16): + gen = model.generate(**enc, max_new_tokens=WINDOW * 2, num_beams=1) + preds.extend(tok.batch_decode(gen, skip_special_tokens=True)) + if (i // 32) % 10 == 0: + print(f"[gen] {i + len(batch)}/{len(all_windows)}", flush=True) + + k = 0 + paragraphs = [] + for text, c in zip(inputs, counts): + stitched = " ".join(preds[k:k + c]) + k += c + paragraphs.append(project_haraqat(stitched, text)) + + csv_path = Path("/tmp/sadeed_r8.csv") + pd.DataFrame({"gt": outputs, "pred": paragraphs}).to_csv(csv_path, index=False, header=False) + + from sadeed_evaluator import ArabicDiacritizationEvaluator as E + + print(f"===== {RUN} SadeedDiac-25 (their default protocol) =====", flush=True) + E.report_errors_on_csv_file(str(csv_path), ground_truth_column_index=0, + predicted_column_index=1, has_header=False, + gt_missing_diacritic_is_error=False) + (CKPT / "sadeed_preds_r8.csv").write_text(csv_path.read_text(), encoding="utf-8") + (CKPT / "EVAL_DONE").write_text("ok\n") + print("[done] EVAL_DONE", flush=True) + + +if __name__ == "__main__": + main() diff --git a/train_hebrew_plane_hf.py b/train_hebrew_plane_hf.py new file mode 100644 index 0000000..1371fb1 --- /dev/null +++ b/train_hebrew_plane_hf.py @@ -0,0 +1,330 @@ +# /// script +# requires-python = ">=3.11" +# dependencies = [ +# "torch==2.5.1", +# "transformers==4.46.3", +# "accelerate>=1.1.0", +# ] +# /// +"""run-022 — Hebrew plane scale-up on HF Jobs (WO04). + +run-021 recipe (byt5-small, 2ep, K=2, 12.48 DER) scaled to byt5-base, +4 epochs, K=3. Data: hf://datasets/Interscript/hebrew-v4 (ro mount). +Checkpoints: hf://buckets/Interscript/isx-training (rw mount). + +Gate (pre-registered, nakdimon test via /data/test.jsonl, greedy, +seq2seq_der, the exact run-021 protocol): DER <= 10.0 (shipped 1.0 is +12.48). Resume contract: labels.sha + step dirs + per-marker eval +progress + original-separator stitching (run-021 lessons intact). + +Usage: + hf jobs uv run --flavor a100-large -d --timeout 24h \ + -v hf://datasets/Interscript/hebrew-v4:/train_data:ro \ + -v hf://buckets/Interscript/isx-training:/ckpt:rw \ + -v /tmp/isx-hf-code:/code:ro \ + train_hebrew_plane_hf.py +""" + +from __future__ import annotations + +import json +import random +import re +import sys +from pathlib import Path + +sys.path.insert(0, "/code") +sys.path.insert(0, "/code/src") + +DATA = Path("/train_data") +CKPT = Path("/ckpt") +RUN = "run-022-heb-plane-base" +UNIT_BYTES = 1400 +N_VAL = 2_000 +EPOCHS = 4 +MAX_LINES = 600_000 +K_PASSES = 3 +MICRO_BS = 8 +GRAD_ACCUM = 2 + +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 rababa.evaluate import seq2seq_der + + +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 + + +def load_jsonl(path: Path) -> list[dict]: + return [json.loads(l) for l in path.read_text(encoding="utf-8").splitlines() if l.strip()] + + +def main() -> None: + run_dir = CKPT / RUN + run_dir.mkdir(parents=True, exist_ok=True) + done_marker = run_dir / "EVAL_DONE" + prog_name = f"eval_progress_{done_marker.name}.jsonl" + + units = [] + for row in load_jsonl(DATA / "train.jsonl"): + t = row["tgt"].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) + + inv_path = run_dir / "inventory.json" + if inv_path.exists(): + classes = json.loads(inv_path.read_text(encoding="utf-8")) + else: + from collections import Counter + + 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)", flush=True) + class_to_id = {c: i for i, c in enumerate(classes)} + n_classes = len(classes) + MASK_ID = n_classes + import hashlib + + labels_digest = hashlib.sha256(inv_path.read_bytes()).hexdigest()[:12] + + tok = AutoTokenizer.from_pretrained("google/byt5-base") + + 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=MICRO_BS, shuffle=True, collate_fn=collate, + num_workers=2, drop_last=True, + ) + total_steps = (len(train_loader) // GRAD_ACCUM) * 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-base") + 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 run_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) + + def save_ckpt(s: int) -> None: + ck = run_dir / f"step-{s}" + 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") + print(f"[ckpt] step-{s} saved", flush=True) + + optimizer.zero_grad() + micro = 0 + 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)) / GRAD_ACCUM + loss.backward() + micro += 1 + if micro % GRAD_ACCUM == 0: + 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) * GRAD_ACCUM:.4f} q={q:.2f}", flush=True) + if step % 1000 == 0: + save_ckpt(step) + + best = run_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") + print("[ckpt] best saved", flush=True) + + # ---- gate: nakdimon test (/data/test.jsonl), K-pass, seq2seq_der ---- + encoder.eval() + plane_emb.eval() + head.eval() + + targets = [r["tgt"] for r in load_jsonl(DATA / "test.jsonl")] + 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 = run_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] + with prog.open("a", encoding="utf-8") as prog_out: + for n_done, i in enumerate(missing, 1): + pred = predict_window(all_windows[i]) + prog_out.write(json.dumps({"i": i, "pred": pred}, ensure_ascii=False) + "\n") + saved[i] = pred + if n_done % 200 == 0: + prog_out.flush() + print(f"[gen] {len(saved)}/{len(all_windows)}", flush=True) + prog.write_text(prog.read_text(encoding="utf-8"), encoding="utf-8") # ensure flushed + + # stitch with ORIGINAL separators (run-021 lesson) + k = 0 + total_wrong = 0.0 + total_positions = 0 + for tgt, c in zip(targets, counts_w): + 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) + d, n = seq2seq_der(pred, tgt) + total_wrong += d * n + total_positions += n + + der = total_wrong / max(1, total_positions) + verdict = {"run": RUN, "steps": step, "k_passes": K_PASSES, "der": der, + "gate": 10.0, "baseline_1_0": 12.48} + (run_dir / "verdict.json").write_text(json.dumps(verdict, indent=1), encoding="utf-8") + print(f"\n===== {RUN} nakdimon greedy DER: {der:.4f} " + f"(gate <= 10.0; heb-diac-plane-1.0 = 12.48) =====", flush=True) + done_marker.touch() + print("[done] EVAL_DONE", flush=True) + + +if __name__ == "__main__": + main() diff --git a/train_thai_tiny_hf.py b/train_thai_tiny_hf.py new file mode 100644 index 0000000..cec6b86 --- /dev/null +++ b/train_thai_tiny_hf.py @@ -0,0 +1,288 @@ +# /// script +# requires-python = ">=3.11" +# dependencies = [ +# "torch==2.5.1", +# "transformers==4.46.3", +# "datasets>=3.0", +# "accelerate>=1.1.0", +# "onnx==1.17.0", +# "onnxruntime==1.20.1", +# "numpy<2", +# ] +# /// +"""WO08: Thai small-tier student — sequence-level KD into a ~12M model. + +Pipeline in one job: stream a Thai Wikipedia slice, teacher-label with +B-K/umt5-thai-g2p-v2-0.5k (the published teacher lineage; NO LLM labels), +train a compact char encoder-decoder (d=256, 4+4 layers) on teacher +outputs, then gate on the kaikki test with the corpus-PER harness. + +Gate (pre-registered): PER <= 3.5% (greedy, corpus-level, true +Levenshtein) AND int8 export < 10 MB. Ship only if both hold; the +artifact contract (kind=tiny-g2p) is designed after the gate. + +Usage: + hf jobs uv run --flavor a10g-large -d --timeout 12h \ + -v hf://datasets/Interscript/thai-g2p:/train_data:ro \ + -v hf://buckets/Interscript/isx-training:/ckpt:rw \ + train_thai_tiny_hf.py +""" + +from __future__ import annotations + +import json +import math +import time +from pathlib import Path + +import torch +import torch.nn as nn +from torch.utils.data import DataLoader, Dataset + +CKPT = Path("/ckpt/run-025-thai-tiny") +DATA = Path("/train_data") +TEACHER = "B-K/umt5-thai-g2p-v2-0.5k" +N_PARAS = 200_000 +MAX_SRC = 300 +D_MODEL = 256 +N_LAYERS = 4 +N_HEADS = 4 +D_FF = 1024 +EPOCHS = 4 +BS = 128 +GATE_PER = 3.5 +GATE_MB = 10.0 +PER = 100.0 + + +def lev(a: str, b: str) -> int: + prev = list(range(len(b) + 1)) + for i, ca in enumerate(a, 1): + cur = [i] + for j, cb in enumerate(b, 1): + cur.append(min(prev[j] + 1, cur[j - 1] + 1, prev[j - 1] + (ca != cb))) + prev = cur + return prev[-1] + + +def corpus_per(pairs: list[tuple[str, str]]) -> float: + dist = sum(lev(p, t) for p, t in pairs) + total = sum(len(t) for _, t in pairs) + return PER * dist / max(1, total) + + +def build_corpus() -> list[tuple[str, str]]: + """teacher-labeled pairs: kaikki test-side srcs + a wiki slice.""" + pairs: list[tuple[str, str]] = [] + import pandas as pd + + from transformers import AutoModelForSeq2SeqLM, AutoTokenizer + + device = "cuda" + tok = AutoTokenizer.from_pretrained(TEACHER) + teacher = AutoModelForSeq2SeqLM.from_pretrained( + TEACHER, torch_dtype=torch.bfloat16).to(device).eval() + + srcs: list[str] = [] + test_path = DATA / "kaikki-test.jsonl" + for line in test_path.read_text(encoding="utf-8").splitlines(): + if line.strip(): + srcs.append(json.loads(line)["src"]) + print(f"[data] kaikki srcs={len(srcs)}", flush=True) + + from datasets import load_dataset + + ds = load_dataset("wikimedia/wikipedia", "20231101.th", split="train", streaming=True) + n = 0 + for row in ds: + for para in row["text"].split("\n"): + p = para.strip() + if 10 <= len(p) <= MAX_SRC: + srcs.append(p) + n += 1 + if n >= N_PARAS: + break + print(f"[data] total srcs={len(srcs)}", flush=True) + + t0 = time.time() + with torch.no_grad(): + for i in range(0, len(srcs), 64): + batch = srcs[i:i + 64] + enc = tok(batch, return_tensors="pt", padding=True, truncation=True, + max_length=MAX_SRC).to(device) + with torch.autocast("cuda", torch.bfloat16): + gen = teacher.generate(**enc, max_new_tokens=MAX_SRC, num_beams=1) + outs = tok.batch_decode(gen, skip_special_tokens=True) + for s, t in zip(batch, outs): + if t.strip(): + pairs.append((s, t)) + if (i // 64) % 100 == 0: + print(f"[teacher] {i + len(batch)}/{len(srcs)} " + f"({(i + len(batch)) / (time.time() - t0):.0f} src/s)", flush=True) + return pairs + + +class TinyG2P(nn.Module): + """Compact char encoder-decoder with a small IPA vocab.""" + + def __init__(self, vocab: dict[str, int]): + super().__init__() + self.vocab = vocab + n_v = len(vocab) + self.emb = nn.Embedding(n_v, D_MODEL) + layer = nn.TransformerEncoderLayer( + D_MODEL, N_HEADS, D_FF, dropout=0.1, batch_first=True, norm_first=True) + self.enc = nn.TransformerEncoder(layer, N_LAYERS) + dlayer = nn.TransformerDecoderLayer( + D_MODEL, N_HEADS, D_FF, dropout=0.1, batch_first=True, norm_first=True) + self.dec = nn.TransformerDecoder(dlayer, N_LAYERS) + self.head = nn.Linear(D_MODEL, n_v) + + def forward(self, src, tgt_in, src_pad, tgt_pad): + mem = self.enc(self.emb(src), src_key_padding_mask=src_pad) + mask = nn.Transformer.generate_square_subsequent_mask( + tgt_in.size(1), device=src.device) + h = self.dec(self.emb(tgt_in), mem, tgt_mask=mask, + memory_key_padding_mask=src_pad, tgt_key_padding_mask=tgt_pad) + return self.head(h) + + +def main() -> None: + CKPT.mkdir(parents=True, exist_ok=True) + pairs_path = CKPT / "distill.jsonl" + if pairs_path.exists(): + pairs = [(json.loads(l)["src"], json.loads(l)["tgt"]) + for l in pairs_path.read_text(encoding="utf-8").splitlines() if l.strip()] + else: + pairs = build_corpus() + with pairs_path.open("w", encoding="utf-8") as f: + for s, t in pairs: + f.write(json.dumps({"src": s, "tgt": t}, ensure_ascii=False) + "\n") + print(f"[data] pairs={len(pairs)}", flush=True) + + vocab = {"": 0, "": 1, "": 2} + for s, t in pairs: + for ch in s + t: + if ch not in vocab: + vocab[ch] = len(vocab) + test = [(json.loads(l)["src"], json.loads(l)["tgt"]) + for l in (DATA / "kaikki-test.jsonl").read_text().splitlines() if l.strip()] + for s, t in test: + for ch in s + t: + if ch not in vocab: + vocab[ch] = len(vocab) + print(f"[vocab] {len(vocab)}", flush=True) + + # kaikki-test srcs must NOT be in training pairs (contamination guard) + test_srcs = {s for s, _ in test} + pairs = [(s, t) for s, t in pairs if s not in test_srcs] + + model = TinyG2P(vocab).cuda().train() + n_params = sum(p.numel() for p in model.parameters()) + print(f"[model] params={n_params / 1e6:.1f}M", flush=True) + + pad, bos, eos = vocab[""], vocab[""], vocab[""] + + class DS(Dataset): + def __getitem__(self, i): + s, t = pairs[i] + return ([vocab[c] for c in s], [bos] + [vocab[c] for c in t], + [vocab[c] for c in t] + [eos]) + + def __len__(self): + return len(pairs) + + def collate(b): + ms = max(len(x[0]) for x in b) + mt = max(len(x[1]) for x in b) + src = torch.full((len(b), ms), pad, dtype=torch.long) + tin = torch.full((len(b), mt), pad, dtype=torch.long) + tout = torch.full((len(b), mt), pad, dtype=torch.long) + for i, (s, ti, to) in enumerate(b): + src[i, :len(s)] = torch.tensor(s) + tin[i, :len(ti)] = torch.tensor(ti) + tout[i, :len(to)] = torch.tensor(to) + return src, tin, tout + + loader = DataLoader(DS(), batch_size=BS, shuffle=True, collate_fn=collate, + num_workers=2, drop_last=True) + opt = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=0.01) + total = len(loader) * EPOCHS + sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, total) + ce = nn.CrossEntropyLoss(ignore_index=pad) + + step = 0 + for _ep in range(EPOCHS): + for src, tin, tout in loader: + src, tin, tout = src.cuda(), tin.cuda(), tout.cuda() + sp = src == pad + tp = tin == pad + with torch.autocast("cuda", torch.bfloat16): + logits = model(src, tin, sp, tp) + loss = ce(logits.reshape(-1, len(vocab)), tout.reshape(-1)) + loss.backward() + torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) + opt.step() + sched.step() + opt.zero_grad() + step += 1 + if step % 200 == 0: + print(f"[train] {step}/{total} loss={float(loss):.4f}", flush=True) + + # ---- greedy decode kaikki-test, corpus PER ---- + model.eval() + preds = [] + with torch.no_grad(): + for i in range(0, len(test), 64): + batch = [s for s, _ in test[i:i + 64]] + ms = max(len(s) for s in batch) + src = torch.full((len(batch), ms), pad, dtype=torch.long) + for r, s in enumerate(batch): + ids = [vocab[c] for c in s] + src[r, :len(ids)] = torch.tensor(ids) + src = src.cuda() + mem = model.enc(model.emb(src), src_key_padding_mask=(src == pad)) + ys = torch.full((len(batch), 1), bos, dtype=torch.long, device="cuda") + done = torch.zeros(len(batch), dtype=torch.bool) + for _ in range(MAX_SRC): + mask = nn.Transformer.generate_square_subsequent_mask( + ys.size(1), device="cuda") + h = model.dec(model.emb(ys), mem, tgt_mask=mask, + memory_key_padding_mask=(src == pad)) + nxt = h[:, -1].argmax(-1) + nxt[done] = pad + ys = torch.cat([ys, nxt.unsqueeze(1)], 1) + done |= nxt == eos + if done.all(): + break + for row in ys[:, 1:].tolist(): + toks = [c for c in row if c not in (pad, eos)] + preds.append("".join(inv[t] for t in toks)) + inv = {v: k for k, v in vocab.items()} + per = corpus_per(list(zip(preds, [t for _, t in test]))) + print(f"[gate] val PER (greedy, corpus) = {per:.4f}% gate<={GATE_PER}%", flush=True) + + # int8 export size check + fp = CKPT / "tiny-fp32.onnx" + src0 = torch.zeros((1, 8), dtype=torch.long) + tin0 = torch.zeros((1, 4), dtype=torch.long) + sp0 = torch.zeros((1, 8), dtype=torch.bool) + tp0 = torch.zeros((1, 4), dtype=torch.bool) + torch.onnx.export(model, (src0, tin0, sp0, tp0), str(fp), + input_names=["src", "tgt_in", "src_pad", "tgt_pad"], + output_names=["logits"], opset_version=17, dynamo=False) + from onnxruntime.quantization import QuantType, quantize_dynamic + + q = CKPT / "tiny-int8.onnx" + quantize_dynamic(str(fp), str(q), weight_type=QuantType.QInt8) + mb = q.stat().st_size / 1e6 + verdict = {"per": round(per, 4), "int8_mb": round(mb, 1), "params_m": round(n_params / 1e6, 1), + "gate_per": GATE_PER, "gate_mb": GATE_MB, + "passed": bool(per <= GATE_PER and mb < GATE_MB)} + (CKPT / "verdict.json").write_text(json.dumps(verdict, indent=1)) + print(json.dumps(verdict), flush=True) + (CKPT / "DONE").write_text("ok\n") + + +if __name__ == "__main__": + main()