feat: HF Jobs substrate — the TODO.sota training stack

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()

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()

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()