diff --git a/src/gpu/distill_specs.yaml b/src/gpu/distill_specs.yaml new file mode 100644 index 0000000..d64597d --- /dev/null +++ b/src/gpu/distill_specs.yaml @@ -0,0 +1,264 @@ +# Distillation specs — data, not code (architecture review C2). +# Loaded by modal_distill.SPECS; one entry per distill/eval arm. +tha-g2p-small: + teacher: B-K/umt5-thai-g2p-v2-0.5k + teacher_is_hub: 'true' + teacher_volume: secryst + student_init: google/byt5-base + train: thai-ipa-expanded/train.jsonl + train_extra: + - thai-ipa/train.jsonl + - thai-ipa/augmented_epitran.jsonl + val: thai-ipa-expanded/val.jsonl + test: thai-ipa-expanded/test.jsonl + eval_test: thai-ipa/test.jsonl + out: secryst_thai_g2p_distill_small/run-004 + mode: sequence + note: umt5 (sentencepiece) teacher -> ByT5-small byte student; +5pp PER gate +tha-g2p-client: + teacher: B-K/umt5-thai-g2p-v2-0.5k + teacher_is_hub: 'true' + teacher_volume: secryst + student_init: google/byt5-small + train: thai-ipa-expanded/train.jsonl + train_extra: + - thai-ipa/train.jsonl + - thai-ipa/augmented_epitran.jsonl + val: thai-ipa-expanded/val.jsonl + test: thai-ipa-expanded/test.jsonl + eval_test: thai-ipa/test.jsonl + out: secryst_thai_g2p_distill_small/run-003 + mode: sequence + note: eval-only spec for tha-g2p-small-1.0 (client tier) +ara-diac-small: + teacher: rababa_arabic_byt5/run-006-morph/best + teacher_volume: rababa + student_init: google/byt5-small + train: r5-units/domain.txt + train_extra: + - r5-units/replay.txt + unit_limits: + - 24000 + - 6000 + max_len: 1450 + label_beams: '1' + out: rababa_arabic_distill_small/run-002 + labels_file: teacher_labels_v2.jsonl + mode: sequence + note: r6 canonical (2.5793 DER); gate <= 3.07 windowed DER-CE +ara-diac-small-pkm: + teacher: rababa_arabic_byt5/run-006-morph/best + teacher_volume: rababa + student_init: google/byt5-small + pkm: + layer_indices: + - -1 + - -2 + - -3 + n_keys: 128 + topk: 32 + train: r5-units/domain.txt + train_extra: + - r5-units/replay.txt + unit_limits: + - 24000 + - 6000 + max_len: 1450 + label_beams: '1' + out: rababa_arabic_distill_small/run-003-pkm + labels_file: teacher_labels_v2.jsonl + labels_complete: 'true' + mode: sequence + note: PKM memory student; gate <= 3.07 windowed DER-CE; verdict vs run-002's 8.259 full-set (closes + >= 1.0pp?) +ara-diac-small-pkm-muon: + teacher: rababa_arabic_byt5/run-006-morph/best + teacher_volume: rababa + student_init: google/byt5-small + pkm: + layer_indices: + - -1 + - -2 + - -3 + n_keys: 128 + topk: 32 + optimizer: muon + muon_lr: '0.01' + train: r5-units/domain.txt + train_extra: + - r5-units/replay.txt + unit_limits: + - 24000 + - 6000 + max_len: 1450 + label_beams: '1' + out: rababa_arabic_distill_small/run-004-pkm-muon + labels_file: teacher_labels_v2.jsonl + labels_complete: 'true' + mode: sequence + note: Muon A/B arm; same gate and verdict rule as run-003-pkm +ara-diac-small-muon: + teacher: rababa_arabic_byt5/run-006-morph/best + teacher_volume: rababa + student_init: google/byt5-small + optimizer: muon + muon_lr: '0.01' + train: r5-units/domain.txt + train_extra: + - r5-units/replay.txt + unit_limits: + - 24000 + - 6000 + max_len: 1450 + label_beams: '1' + out: rababa_arabic_distill_small/run-005-muon + labels_file: teacher_labels_v2.jsonl + labels_complete: 'true' + mode: sequence + note: vanilla ByT5-small + Muon (factorial cell 4) +ara-diac-small-2: + teacher: rababa_arabic_byt5/run-007-news/best + teacher_volume: rababa + student_init: google/byt5-small + optimizer: muon + muon_lr: '0.01' + train: r5-units/domain.txt + train_extra: + - r5-units/replay.txt + unit_limits: + - 24000 + - 6000 + max_len: 1450 + label_beams: '1' + out: rababa_arabic_distill_small/run-006-r7-muon + labels_file: teacher_labels_r7.jsonl + mode: sequence + note: 'r7 teacher + Muon; E4 gate: beat the shipped 8.259 by >= 2pp' +ara-diac-tiny: + teacher: rababa_arabic_byt5/run-006-morph/best + teacher_volume: rababa + out_volume: secryst + student_config: + d_model: 384 + d_ff: 1536 + num_heads: 6 + enc_layers: 8 + dec_layers: 8 + train: r5-units/domain.txt + train_extra: + - r5-units/replay.txt + unit_limits: + - 8000 + - 4000 + max_len: 1450 + label_beams: '1' + out: rababa_arabic_distill_tiny/run-005 + labels_file: teacher_labels_v2.jsonl + labels_complete: 'true' + mode: sequence + note: clean-label tiny re-run; r6 teacher (2.5793 DER); gate <= 3.07 +tha-g2p-tiny: + teacher: B-K/umt5-thai-g2p-v2-0.5k + teacher_is_hub: 'true' + teacher_volume: secryst + student_config: + d_model: 384 + d_ff: 1536 + num_heads: 6 + enc_layers: 8 + dec_layers: 8 + train: thai-ipa-expanded/train.jsonl + train_extra: + - thai-ipa/train.jsonl + - thai-ipa/augmented_epitran.jsonl + eval_test: thai-ipa/test.jsonl + out: secryst_thai_g2p_distill_tiny/run-001 + labels_complete: 'true' + mode: sequence + label_beams: '4' + max_len: 384 + note: client tier; collapsed from-scratch (75.8 PER) — see run-002 mk +tha-g2p-tiny-mk: + teacher: B-K/umt5-thai-g2p-v2-0.5k + teacher_is_hub: 'true' + teacher_volume: secryst + student_config: + d_model: 384 + d_ff: 1536 + num_heads: 6 + enc_layers: 8 + dec_layers: 8 + eval_test: thai-ipa/test.jsonl + out: secryst_thai_g2p_distill_tiny/run-002 + labels_complete: 'true' + max_len: 384 + note: microkimi bridges; 71.12 PER — improved, not rescued +tha-g2p-mid-mk: + teacher: B-K/umt5-thai-g2p-v2-0.5k + teacher_is_hub: 'true' + teacher_volume: secryst + student_config: + d_model: 512 + d_ff: 2048 + num_heads: 8 + enc_layers: 10 + dec_layers: 10 + eval_test: thai-ipa/test.jsonl + out: secryst_thai_g2p_distill_mid/run-001 + labels_complete: 'true' + max_len: 384 + note: 70M bridge rung; labels reused from tiny run-002 +fas-g2p-tiny: + teacher: persian_g2p/run-001/best + teacher_volume: persian + student_config: + d_model: 384 + d_ff: 1536 + num_heads: 6 + enc_layers: 8 + dec_layers: 8 + data_volume: /datasets + train: persian_g2p/train.jsonl + unit_limits: + - 60000 + test: persian_g2p/test.jsonl + out: interscript_fas_g2p_distill_tiny/run-001 + mode: sequence + label_beams: '4' + note: client tier; collapsed from-scratch (77.97 PER) +heb-diac-tiny: + teacher: rababa_hebrew_byt5_s43/run-001/best + teacher_volume: rababa + out_volume: secryst + student_config: + d_model: 384 + d_ff: 1536 + num_heads: 6 + enc_layers: 8 + dec_layers: 8 + train: hebrew-v4/train.jsonl + out: rababa_hebrew_distill_tiny/run-002 + mode: sequence + label_beams: '4' + note: client tier; collapsed from-scratch (100 DER) +fas-g2p-small: + teacher: persian_g2p/run-001/best + teacher_volume: persian + student_init: google/byt5-small + train: persian_g2p/train.jsonl + val: persian_g2p/val.jsonl + test: persian_g2p/test.jsonl + out: interscript_fas_g2p_distill_small/run-001 + note: ByT5-small teacher (already byte-level) -> ByT5-small student; CER gate +heb-diac-small: + teacher: rababa_hebrew_byt5_s43/run-001/best + student_init: google/byt5-small + train: hebrew-v4/train.jsonl + val: hebrew-v4/val.jsonl + out: rababa_hebrew_distill_small/run-001 +heb-diac-small-s46: + teacher: rababa_hebrew/run-s46-phonikud-plus/run-002-gold-ft/best + student_init: google/byt5-small + train: hebrew-v4/train.jsonl + val: hebrew-v4/val.jsonl + out: rababa_hebrew_distill_small/run-002-s46 diff --git a/src/gpu/modal_distill.py b/src/gpu/modal_distill.py index 6ceece5..4229d84 100644 --- a/src/gpu/modal_distill.py +++ b/src/gpu/modal_distill.py @@ -80,264 +80,23 @@ def _ensure_src_path() -> None: PERSIAN_CHECKPOINTS = modal.Volume.from_name("persian-g2p-checkpoints") PERSIAN_DATASETS = modal.Volume.from_name("persian-g2p-datasets") -SPECS: dict[str, dict[str, str]] = { - "tha-g2p-small": { - # secryst's saved umt5 artifacts are unusable (5.15 dropped the - # untied lm_head) and volume epitran data is tone-less (see - # modal_teacher_thai.py notes) — teacher is the B-K hub base - # directly: 6.37% PER published, verified loading + exact-match - # generations under transformers 5.14.1 (probes 2026-08-18) - "teacher": "B-K/umt5-thai-g2p-v2-0.5k", - "teacher_is_hub": "true", - "teacher_volume": "secryst", - # run-004: ByT5-base student — the small student hit a 12.06% - # generalization ceiling (+7.6pp over teacher, rejected); base - # has 4x capacity at a ~1.2GB artifact, still client-tier - "student_init": "google/byt5-base", - "train": "thai-ipa-expanded/train.jsonl", - "train_extra": ["thai-ipa/train.jsonl", "thai-ipa/augmented_epitran.jsonl"], - "val": "thai-ipa-expanded/val.jsonl", - "test": "thai-ipa-expanded/test.jsonl", - "eval_test": "thai-ipa/test.jsonl", - "out": "secryst_thai_g2p_distill_small/run-004", - "mode": "sequence", # cross-tokenizer: teacher generates, student trains CE - "note": "umt5 (sentencepiece) teacher -> ByT5-small byte student; +5pp PER gate", - }, - "tha-g2p-client": { - # the shipped client rung: run-003, ByT5-small on the full label - # set — capacity-limited (+7.6pp) but the smallest artifact that - # does not collapse (see docs/RESULTS.md frontier table) - "teacher": "B-K/umt5-thai-g2p-v2-0.5k", - "teacher_is_hub": "true", - "teacher_volume": "secryst", - "student_init": "google/byt5-small", - "train": "thai-ipa-expanded/train.jsonl", - "train_extra": ["thai-ipa/train.jsonl", "thai-ipa/augmented_epitran.jsonl"], - "val": "thai-ipa-expanded/val.jsonl", - "test": "thai-ipa-expanded/test.jsonl", - "eval_test": "thai-ipa/test.jsonl", - "out": "secryst_thai_g2p_distill_small/run-003", - "mode": "sequence", - "note": "eval-only spec for tha-g2p-small-1.0 (client tier)", - }, - "ara-diac-small": { - # r5 paragraph-context teacher (2.68 DER-CE windowed @1400B, - # RELEASE-FROZEN) -> ByT5-small student. Contract decode is - # GREEDY with generation cap 2x window (eval_sadeed_windowed). - # Corpus: r5-units joined paragraph units (src = stripped - # diacritics, teacher regenerates the labels). - "teacher": "rababa_arabic_byt5/run-006-morph/best", - "teacher_volume": "rababa", - "student_init": "google/byt5-small", - "train": "r5-units/domain.txt", - "train_extra": ["r5-units/replay.txt"], - "unit_limits": [24000, 6000], - "max_len": 1450, - "label_beams": "1", - "out": "rababa_arabic_distill_small/run-002", - # v2: every label generated before the byt5 decode_joined fix is - # mojibake (double-encoded); relabel from scratch on the new file - "labels_file": "teacher_labels_v2.jsonl", - "mode": "sequence", - "note": "r6 canonical (2.5793 DER); gate <= 3.07 windowed DER-CE", - }, - "ara-diac-small-pkm": { - # TODO.qwen-next/02 — the LongCat/Qwen capacity axis: keep the - # ByT5-small compute, add product-key lookup memory (+~25M - # params). Everything else identical to run-002 (teacher, corpus, - # labels, seed) so the comparison is single-variable. - "teacher": "rababa_arabic_byt5/run-006-morph/best", - "teacher_volume": "rababa", - "student_init": "google/byt5-small", - # byt5-small has only 4 decoder blocks (depth lives in the - # encoder) — all but the first carry memory - "pkm": {"layer_indices": [-1, -2, -3], "n_keys": 128, "topk": 32}, - "train": "r5-units/domain.txt", - "train_extra": ["r5-units/replay.txt"], - "unit_limits": [24000, 6000], - "max_len": 1450, - "label_beams": "1", - "out": "rababa_arabic_distill_small/run-003-pkm", - "labels_file": "teacher_labels_v2.jsonl", - "labels_complete": "true", - "mode": "sequence", - "note": "PKM memory student; gate <= 3.07 windowed DER-CE; " - "verdict vs run-002's 8.259 full-set (closes >= 1.0pp?)", - }, - "ara-diac-small-pkm-muon": { - # TODO.qwen-next/03 — identical to ara-diac-small-pkm except the - # optimizer (Muon + AdamW groups). The A/B pair is run-003-pkm - # (AdamW) vs run-004-pkm-muon. - "teacher": "rababa_arabic_byt5/run-006-morph/best", - "teacher_volume": "rababa", - "student_init": "google/byt5-small", - "pkm": {"layer_indices": [-1, -2, -3], "n_keys": 128, "topk": 32}, - "optimizer": "muon", - "muon_lr": "0.01", - "train": "r5-units/domain.txt", - "train_extra": ["r5-units/replay.txt"], - "unit_limits": [24000, 6000], - "max_len": 1450, - "label_beams": "1", - "out": "rababa_arabic_distill_small/run-004-pkm-muon", - "labels_file": "teacher_labels_v2.jsonl", - "labels_complete": "true", - "mode": "sequence", - "note": "Muon A/B arm; same gate and verdict rule as run-003-pkm", - }, - "ara-diac-small-muon": { - # factorial completion (EXPERIMENTS.md E3 caveat): the A/B was - # PKM+Muon vs PKM+AdamW; this cell measures Muon WITHOUT memory - # so the 2x2 {vanilla, pkm} x {adamw, muon} closes cleanly - "teacher": "rababa_arabic_byt5/run-006-morph/best", - "teacher_volume": "rababa", - "student_init": "google/byt5-small", - "optimizer": "muon", - "muon_lr": "0.01", - "train": "r5-units/domain.txt", - "train_extra": ["r5-units/replay.txt"], - "unit_limits": [24000, 6000], - "max_len": 1450, - "label_beams": "1", - "out": "rababa_arabic_distill_small/run-005-muon", - "labels_file": "teacher_labels_v2.jsonl", - "labels_complete": "true", - "mode": "sequence", - "note": "vanilla ByT5-small + Muon (factorial cell 4)", - }, - "ara-diac-small-2": { - # ara-diac-small-2.0 candidate: the r7 canonical teacher + the - # E3-adopted Muon optimizer. Fresh labels from r7 (the v2 labels - # are r6-teacher). EXPERIMENTS.md E4. - "teacher": "rababa_arabic_byt5/run-007-news/best", - "teacher_volume": "rababa", - "student_init": "google/byt5-small", - "optimizer": "muon", - "muon_lr": "0.01", - "train": "r5-units/domain.txt", - "train_extra": ["r5-units/replay.txt"], - "unit_limits": [24000, 6000], - "max_len": 1450, - "label_beams": "1", - "out": "rababa_arabic_distill_small/run-006-r7-muon", - "labels_file": "teacher_labels_r7.jsonl", - "mode": "sequence", - "note": "r7 teacher + Muon; E4 gate: beat the shipped 8.259 by >= 2pp", - }, - "ara-diac-tiny": { - "teacher": "rababa_arabic_byt5/run-006-morph/best", - "teacher_volume": "rababa", - "out_volume": "secryst", - "student_config": {"d_model": 384, "d_ff": 1536, "num_heads": 6, - "enc_layers": 8, "dec_layers": 8}, - "train": "r5-units/domain.txt", - "train_extra": ["r5-units/replay.txt"], - "unit_limits": [8000, 4000], - "max_len": 1450, - "label_beams": "1", - "out": "rababa_arabic_distill_tiny/run-005", - # run-004 used the Aug-23 snapshot = PRE decode_joined fix: every - # label mojibake (verified: its best/ decodes as UTF-8-as-Latin1, - # and its final_eval.json reproduces the retracted 83.08 exactly). - # run-005 is the actual clean-label re-run the retraction calls for. - "labels_file": "teacher_labels_v2.jsonl", - "labels_complete": "true", - "mode": "sequence", - "note": "clean-label tiny re-run; r6 teacher (2.5793 DER); gate <= 3.07", - }, - "tha-g2p-tiny": { - "teacher": "B-K/umt5-thai-g2p-v2-0.5k", - "teacher_is_hub": "true", - "teacher_volume": "secryst", - "student_config": {"d_model": 384, "d_ff": 1536, "num_heads": 6, - "enc_layers": 8, "dec_layers": 8}, - "train": "thai-ipa-expanded/train.jsonl", - "train_extra": ["thai-ipa/train.jsonl", "thai-ipa/augmented_epitran.jsonl"], - "eval_test": "thai-ipa/test.jsonl", - "out": "secryst_thai_g2p_distill_tiny/run-001", - "labels_complete": "true", - "mode": "sequence", - "label_beams": "4", - "max_len": 384, - "note": "client tier; collapsed from-scratch (75.8 PER) — see run-002 mk", - }, - "tha-g2p-tiny-mk": { - "teacher": "B-K/umt5-thai-g2p-v2-0.5k", - "teacher_is_hub": "true", - "teacher_volume": "secryst", - "student_config": {"d_model": 384, "d_ff": 1536, "num_heads": 6, - "enc_layers": 8, "dec_layers": 8}, - "eval_test": "thai-ipa/test.jsonl", - "out": "secryst_thai_g2p_distill_tiny/run-002", - "labels_complete": "true", - "max_len": 384, - "note": "microkimi bridges; 71.12 PER — improved, not rescued", - }, - "tha-g2p-mid-mk": { - "teacher": "B-K/umt5-thai-g2p-v2-0.5k", - "teacher_is_hub": "true", - "teacher_volume": "secryst", - "student_config": {"d_model": 512, "d_ff": 2048, "num_heads": 8, - "enc_layers": 10, "dec_layers": 10}, - "eval_test": "thai-ipa/test.jsonl", - "out": "secryst_thai_g2p_distill_mid/run-001", - "labels_complete": "true", - "max_len": 384, - "note": "70M bridge rung; labels reused from tiny run-002", - }, - "fas-g2p-tiny": { - "teacher": "persian_g2p/run-001/best", - "teacher_volume": "persian", - "student_config": {"d_model": 384, "d_ff": 1536, "num_heads": 6, - "enc_layers": 8, "dec_layers": 8}, - "data_volume": "/datasets", - "train": "persian_g2p/train.jsonl", - "unit_limits": [60000], - "test": "persian_g2p/test.jsonl", - "out": "interscript_fas_g2p_distill_tiny/run-001", - "mode": "sequence", - "label_beams": "4", - "note": "client tier; collapsed from-scratch (77.97 PER)", - }, - "heb-diac-tiny": { - "teacher": "rababa_hebrew_byt5_s43/run-001/best", - "teacher_volume": "rababa", - "out_volume": "secryst", - "student_config": {"d_model": 384, "d_ff": 1536, "num_heads": 6, - "enc_layers": 8, "dec_layers": 8}, - "train": "hebrew-v4/train.jsonl", - "out": "rababa_hebrew_distill_tiny/run-002", - "mode": "sequence", - "label_beams": "4", - "note": "client tier; collapsed from-scratch (100 DER)", - }, - "fas-g2p-small": { - "teacher": "persian_g2p/run-001/best", - "teacher_volume": "persian", - "student_init": "google/byt5-small", - "train": "persian_g2p/train.jsonl", - "val": "persian_g2p/val.jsonl", - "test": "persian_g2p/test.jsonl", - "out": "interscript_fas_g2p_distill_small/run-001", - "note": "ByT5-small teacher (already byte-level) -> ByT5-small student; CER gate", - }, - "heb-diac-small": { - "teacher": "rababa_hebrew_byt5_s43/run-001/best", - "student_init": "google/byt5-small", - "train": "hebrew-v4/train.jsonl", - "val": "hebrew-v4/val.jsonl", - "out": "rababa_hebrew_distill_small/run-001", - }, - "heb-diac-small-s46": { - # re-distill the client tier from the s46 teacher (greedy 16.44 - # vs s43's 29.0) — the 1.0 student tracked its teacher's greedy - "teacher": "rababa_hebrew/run-s46-phonikud-plus/run-002-gold-ft/best", - "student_init": "google/byt5-small", - "train": "hebrew-v4/train.jsonl", - "val": "hebrew-v4/val.jsonl", - "out": "rababa_hebrew_distill_small/run-002-s46", - }, -} +def _load_specs() -> dict: + """SPECS as data (distill_specs.yaml) — the entry file is copied to + /root/.py on Modal while the repo image sits at + /root/interscript-ml, so try both layouts.""" + import yaml + + for cand in ( + Path(__file__).resolve().parent / "distill_specs.yaml", + Path.cwd() / "src/gpu/distill_specs.yaml", + Path("/root/interscript-ml/src/gpu/distill_specs.yaml"), + ): + if cand.exists(): + return yaml.safe_load(cand.read_text(encoding="utf-8")) + raise RuntimeError("distill_specs.yaml not found on any known layout") + + +SPECS: dict[str, dict[str, str]] = _load_specs() app = modal.App("interscript-ml-distill", image=IMAGE) @@ -1228,14 +987,17 @@ def evaluate_der(spec_id: str, window: int = 1400, limit: int = 0) -> dict: cap, stitch, project haraqat onto the input letters (zero-skip), DER-CE via the Misraj evaluator. Gate: teacher + 0.5pp (DISTILL-SOURCE-PROMPT: 3.18 target from the 2.68 teacher).""" - import difflib - import re from pathlib import Path import pyarrow.parquet as pq - import torch from transformers import AutoModelForSeq2SeqLM, AutoTokenizer + _ensure_src_path() + from harness.sadeed import ( + strip_diacritics, + windowed_paragraphs, + ) + spec = SPECS[spec_id] vol_map = { "rababa": "/checkpoints", @@ -1263,69 +1025,14 @@ def evaluate_der(spec_id: str, window: int = 1400, limit: int = 0) -> dict: for m in (teacher, student): m.generation_config.max_length = 100_000 - diac = re.compile("[ً-ٰٟۖ-ۭ]") table = pq.read_table("/opt/rababa/data/sadeed-diac-25/train.parquet") - inputs = [diac.sub("", t) for t in table.column("input").to_pylist()] + inputs = [strip_diacritics(t) for t in table.column("input").to_pylist()] gts = table.column("output").to_pylist() if limit: inputs, gts = inputs[:limit], gts[:limit] - def split_windows(text: str) -> list[str]: - if len(text.encode()) <= window: - return [text] - wins, cur, n = [], [], 0 - for w in text.split(): - c = len(w.encode()) + 1 - if cur and n + c > window: - 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: - haraqat = [""] - for ch in pred: - if diac.match(ch): - haraqat[-1] += ch - else: - haraqat.append("") - haraqat = haraqat[1:] - pred_letters = [c for c in pred if not diac.match(c)] - text_letters = [c for c in text if not diac.match(c)] - sm = difflib.SequenceMatcher(None, text_letters, pred_letters, autojunk=False) - out = [] - for op, i1, i2, j1, _ in sm.get_opcodes(): - if op == "equal": - for k in range(i2 - i1): - out.append(text_letters[i1 + k] + haraqat[j1 + k]) - else: - for k in range(i1, i2): - out.append(text_letters[k]) - return "".join(out) - def der_ce(model) -> dict: - windows, counts = [], [] - for text in inputs: - ws = split_windows(text) - counts.append(len(ws)) - windows.extend(ws) - preds = [] - with torch.no_grad(): - for i in range(0, len(windows), 8): - batch = windows[i : i + 8] - 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)) - k = 0 - paragraphs = [] - for text, c in zip(inputs, counts, strict=True): - paragraphs.append(project_haraqat(" ".join(preds[k : k + c]), text)) - k += c + paragraphs = windowed_paragraphs(model, tok, inputs, window=window) import sys diff --git a/src/gpu/modal_export.py b/src/gpu/modal_export.py index cdefffb..92d7286 100644 --- a/src/gpu/modal_export.py +++ b/src/gpu/modal_export.py @@ -270,7 +270,21 @@ def parity_model(model_id: str, precisions: list[str], limit: int = 0) -> dict[s if limit: pairs = pairs[:limit] + def stage(event: str) -> None: + # the gate's observable interface: durable stage log on the + # volume (silent container kills are otherwise undiagnosable — + # seven consecutive ara-diac2 attempts died without a traceback) + import time as _time + + out_dir0 = Path("/outputs/imf") / model_id + out_dir0.mkdir(parents=True, exist_ok=True) + with (out_dir0 / "parity_stages.jsonl").open("a", encoding="utf-8") as fh: + fh.write(f'{{"t": {round(_time.time())}, "event": "{event}"}}\n') + MODELS_VOLUME.commit() + + stage(f"start model={model_id} pairs={len(pairs)}") reference = reference_decode(model, [src for src, _ in pairs], max_len=128) + stage("reference-decode done") out_dir = Path("/outputs/imf") / model_id meta_path = Path("/root/interscript-ml", spec["metadata"]) @@ -278,7 +292,9 @@ def parity_model(model_id: str, precisions: list[str], limit: int = 0) -> dict[s reports: dict[str, str] = {} for precision in precisions: zip_path = out_dir / f"{mid}-{precision}.zip" + stage(f"onnx decode {precision}") report = run_parity(model, zip_path, pairs, max_len=128, reference=reference) + stage(f"parity report {precision} delta={report.cer_delta}") reports[precision] = ( f"samples={report.samples} cer_ref={report.cer_reference}pp " f"cer_onnx={report.cer_onnx}pp delta={report.cer_delta}pp " @@ -288,6 +304,7 @@ def parity_model(model_id: str, precisions: list[str], limit: int = 0) -> dict[s raise RuntimeError(f"parity gate FAILED for {zip_path.name}") write_parity(zip_path, report) margins = run_margin_analysis(model, zip_path, pairs, max_len=128) + stage(f"margin report {precision} flips={margins.flip_rate:.4%}") write_margin_report(margins, out_dir / f"{mid}-margins-{precision}.json") reports[precision] += ( f" | margin flips={margins.flip_rate:.4%} kld={margins.kld_mean:.2e} " @@ -375,184 +392,3 @@ def margins(model: str, precisions: str = "fp32,fp16,int8", limit: int = 0) -> N reports = margin_model.remote(model, precisions.split(","), limit) for precision, status in reports.items(): print(f"{model} [{precision}] {status}") - - -@app.function( - cpu=8, - memory=32 * 1024, - timeout=5 * 3600, - volumes={**CHECKPOINT_VOLUMES, **DATASET_VOLUMES, "/outputs": MODELS_VOLUME}, -) -def int8_pc_probe(model_id: str = "heb-diac", limit: int = 300) -> dict: - """E1 follow-up: does per-channel int8 remove the confident-position - argmax flips? Rebuilds the int8 graphs from the fp32 zip with - per_channel=True, packages them as a probe zip (copy of the shipped - int8 zip with graphs swapped — NOT a release artifact), and compares - margin reports on the same pairs.""" - import sys - import tempfile - import zipfile - - sys.path.insert(0, "/root/interscript-ml/src") - - spec = MODELS[model_id] - checkpoint = Path(spec["volume"]) / spec["checkpoint"] - test_path = Path(spec["test_volume"]) / spec["test_data"] - - from imf.export import load_byte_seq2seq, quantize_int8 - from imf.parity import run_margin_analysis - - model = load_byte_seq2seq(checkpoint) - pairs = _load_pairs(test_path)[:limit] - - out_dir = Path("/outputs/imf") / model_id - meta_path = Path("/root/interscript-ml", spec["metadata"]) - mid = re.search(r"^id:\s*(\S+)", meta_path.read_text(encoding="utf-8"), re.M).group(1) - fp32_zip = out_dir / f"{mid}-fp32.zip" - int8_zip = out_dir / f"{mid}-int8.zip" - if not fp32_zip.exists() or not int8_zip.exists(): - raise RuntimeError(f"need both {fp32_zip.name} and {int8_zip.name} on the volume") - - shipped = run_margin_analysis(model, int8_zip, pairs, max_len=128) - - with tempfile.TemporaryDirectory() as tmp: - tmp = Path(tmp) - with zipfile.ZipFile(fp32_zip) as zf: - zf.extract("encoder.onnx", tmp) - dec = "decoder-kv.onnx" if "decoder-kv.onnx" in zf.namelist() else "decoder.onnx" - zf.extract(dec, tmp) - enc_pc = tmp / "encoder-pc.onnx" - dec_pc = tmp / dec.replace(".onnx", "-pc.onnx") - quantize_int8(tmp / "encoder.onnx", enc_pc, per_channel=True) - quantize_int8(tmp / dec, dec_pc, per_channel=True) - - probe_zip = tmp / f"{mid}-int8-pc-probe.zip" - with zipfile.ZipFile(int8_zip) as src, zipfile.ZipFile( - probe_zip, "w", zipfile.ZIP_DEFLATED - ) as dst: - for name in src.namelist(): - if name == "encoder.onnx": - dst.writestr(name, enc_pc.read_bytes()) - elif name == dec: - dst.writestr(name, dec_pc.read_bytes()) - else: - dst.writestr(name, src.read(name)) - - per_channel_report = run_margin_analysis(model, probe_zip, pairs, max_len=128) - size_shipped = int8_zip.stat().st_size - size_probe = probe_zip.stat().st_size - - def row(r): - return { - "flips": r.flipped_tokens, "tokens": r.tokens, - "flip_rate": r.flip_rate, "kld_mean": r.kld_mean, - "flip_low_margin_share": r.flip_low_margin_share, - } - - return { - "model": model_id, "pairs": len(pairs), - "shipped_int8": row(shipped), "per_channel_int8": row(per_channel_report), - "size_bytes": {"shipped": size_shipped, "per_channel": size_probe}, - } - - -@app.local_entrypoint() -def int8_pc(model: str = "heb-diac", limit: int = 300) -> None: - print(int8_pc_probe.remote(model, limit)) - - -@app.function( - cpu=8, - memory=32 * 1024, - timeout=5 * 3600, - volumes={**CHECKPOINT_VOLUMES, **DATASET_VOLUMES, "/outputs": MODELS_VOLUME}, -) -def int8_head_probe(model_id: str = "heb-diac", limit: int = 300) -> dict: - """E1 follow-up 2: per-channel alone did NOT fix heb-diac's 9.3% - flip rate (8.5% remaining, 78% still at confident positions). This - probe keeps the logits-producing MatMul (the tied head) in fp32 and - quantizes only the body — per-tensor and per-channel variants.""" - import sys - import tempfile - import zipfile - - sys.path.insert(0, "/root/interscript-ml/src") - - spec = MODELS[model_id] - checkpoint = Path(spec["volume"]) / spec["checkpoint"] - test_path = Path(spec["test_volume"]) / spec["test_data"] - - from imf.export import head_matmul_names, load_byte_seq2seq - from imf.parity import run_margin_analysis - - model = load_byte_seq2seq(checkpoint) - pairs = _load_pairs(test_path)[:limit] - - out_dir = Path("/outputs/imf") / model_id - meta_path = Path("/root/interscript-ml", spec["metadata"]) - mid = re.search(r"^id:\s*(\S+)", meta_path.read_text(encoding="utf-8"), re.M).group(1) - fp32_zip = out_dir / f"{mid}-fp32.zip" - int8_zip = out_dir / f"{mid}-int8.zip" - if not fp32_zip.exists() or not int8_zip.exists(): - raise RuntimeError(f"need both {fp32_zip.name} and {int8_zip.name} on the volume") - - shipped = run_margin_analysis(model, int8_zip, pairs, max_len=128) - - results: dict = { - "model": model_id, "pairs": len(pairs), - "shipped_int8": { - "flips": shipped.flipped_tokens, "tokens": shipped.tokens, - "flip_rate": shipped.flip_rate, "kld_mean": shipped.kld_mean, - "flip_low_margin_share": shipped.flip_low_margin_share, - }, - } - - with tempfile.TemporaryDirectory() as tmp: - tmp = Path(tmp) - with zipfile.ZipFile(fp32_zip) as zf: - zf.extract("encoder.onnx", tmp) - dec = "decoder-kv.onnx" if "decoder-kv.onnx" in zf.namelist() else "decoder.onnx" - zf.extract(dec, tmp) - - head_nodes = head_matmul_names(tmp / dec) - results["head_nodes_excluded"] = head_nodes - - from onnxruntime.quantization import QuantType, quantize_dynamic - - for variant, per_channel in (("head32", False), ("head32_pc", True)): - enc_q = tmp / f"encoder-{variant}.onnx" - dec_q = tmp / dec.replace(".onnx", f"-{variant}.onnx") - quantize_dynamic( - str(tmp / "encoder.onnx"), str(enc_q), - weight_type=QuantType.QInt8, op_types_to_quantize=["MatMul"], - per_channel=per_channel, - ) - quantize_dynamic( - str(tmp / dec), str(dec_q), - weight_type=QuantType.QInt8, op_types_to_quantize=["MatMul"], - per_channel=per_channel, nodes_to_exclude=head_nodes, - ) - probe_zip = tmp / f"{mid}-int8-{variant}-probe.zip" - with zipfile.ZipFile(int8_zip) as src, zipfile.ZipFile( - probe_zip, "w", zipfile.ZIP_DEFLATED - ) as dst: - for name in src.namelist(): - if name == "encoder.onnx": - dst.writestr(name, enc_q.read_bytes()) - elif name == dec: - dst.writestr(name, dec_q.read_bytes()) - else: - dst.writestr(name, src.read(name)) - report = run_margin_analysis(model, probe_zip, pairs, max_len=128) - results[f"int8_{variant}"] = { - "flips": report.flipped_tokens, "tokens": report.tokens, - "flip_rate": report.flip_rate, "kld_mean": report.kld_mean, - "flip_low_margin_share": report.flip_low_margin_share, - "size_bytes": probe_zip.stat().st_size, - } - return results - - -@app.local_entrypoint() -def int8_head(model: str = "heb-diac", limit: int = 300) -> None: - print(int8_head_probe.remote(model, limit)) diff --git a/src/harness/sadeed.py b/src/harness/sadeed.py new file mode 100644 index 0000000..1b74796 --- /dev/null +++ b/src/harness/sadeed.py @@ -0,0 +1,106 @@ +"""The Sadeed windowed harness — the campaign's central measurement +protocol, in one place. + +Windowed zero-skip (the published protocol): inputs over the byte +budget split at word boundaries; greedy decode with a 2x-window +generation cap (diacritized output runs 1.4-1.6x input — a shorter cap +silently truncates the hardest paragraphs, which the evaluator then +skips: survivorship bias, not quality); window predictions stitched and +haraqat projected onto the input letters so output structure always +matches ground truth. + +Previously pasted verbatim into rababa train_arabic_r5/r6/r7/r8 and +modal_distill.evaluate_der — five copies of a protocol every published +Arabic number depends on. rababa scripts should vendor this single +module instead of carrying inline copies. +""" + +from __future__ import annotations + +import re + +DIACRITICS_RE = re.compile("[ؐ-ًؚ-ٰٟۖ-ۜ۟-۪ۨ-ۭ]") + + +def split_windows(text: str, budget: int = 1400) -> list[str]: + """Split at word boundaries so no window exceeds the byte budget.""" + if len(text.encode("utf-8")) <= budget: + return [text] + windows: list[str] = [] + current: list[str] = [] + n = 0 + for word in text.split(): + cost = len(word.encode("utf-8")) + 1 + if current and n + cost > budget: + windows.append(" ".join(current)) + current, n = [], 0 + current.append(word) + n += cost + if current: + windows.append(" ".join(current)) + return windows + + +def project_haraqat(pred: str, text: str) -> str: + """Project predicted haraqat onto the input's letters, so the output + structure matches ground truth even when the prediction inserts or + drops letters (zero-skip contract).""" + 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: list[str] = [] + 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 strip_diacritics(text: str) -> str: + return DIACRITICS_RE.sub("", text) + + +def windowed_paragraphs(model, tokenizer, inputs, window: int = 1400, + batch_size: int = 8, device: str = "cuda") -> list[str]: + """The full protocol over a paragraph list: split, greedy-decode each + window (2x-window cap) under bf16 autocast, stitch, project. Returns + haraqat-projected paragraphs ready for the Misraj evaluator.""" + import torch + + windows: list[str] = [] + counts: list[int] = [] + for text in inputs: + ws = split_windows(text, window) + counts.append(len(ws)) + windows.extend(ws) + + preds: list[str] = [] + with torch.no_grad(): + for i in range(0, len(windows), batch_size): + batch = windows[i : i + batch_size] + enc = tokenizer( + batch, return_tensors="pt", padding=True, truncation=True, + max_length=window, + ).to(device) + with torch.autocast("cuda", torch.bfloat16): + gen = model.generate(**enc, max_new_tokens=window * 2, num_beams=1) + preds.extend(tokenizer.batch_decode(gen, skip_special_tokens=True)) + + paragraphs: list[str] = [] + k = 0 + for text, c in zip(inputs, counts, strict=True): + paragraphs.append(project_haraqat(" ".join(preds[k : k + c]), text)) + k += c + return paragraphs diff --git a/tests/test_harness_sadeed.py b/tests/test_harness_sadeed.py new file mode 100644 index 0000000..23b4f0c --- /dev/null +++ b/tests/test_harness_sadeed.py @@ -0,0 +1,65 @@ +"""Unit tests for the Sadeed windowed harness — the published protocol, +tested directly instead of through 6-hour GPU runs.""" + +from __future__ import annotations + +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "src")) + +from harness.sadeed import ( # noqa: E402 + DIACRITICS_RE, + project_haraqat, + split_windows, + strip_diacritics, +) + + +def test_split_windows_short_text_is_single_window() -> None: + assert split_windows("قصيرة", 1400) == ["قصيرة"] + + +def test_split_windows_respects_byte_budget() -> None: + words = ["كلمة"] * 50 # 2 bytes header + 4 bytes per word + space + text = " ".join(words) + budget = 60 + windows = split_windows(text, budget) + assert len(windows) > 1 + for w in windows: + assert len(w.encode("utf-8")) <= budget + + +def test_split_windows_never_drops_words() -> None: + words = ["a" * 30] * 7 + text = " ".join(words) + windows = split_windows(text, 100) + rejoined = " ".join(windows).split() + assert rejoined == text.split() + + +def test_project_haraqat_identity_on_equal_letters() -> None: + text = strip_diacritics("قَالَ") + assert project_haraqat("قَالَ", text) == "قَالَ" + + +def test_project_haraqat_survives_prediction_insertions() -> None: + # prediction drops a letter (ب) — projection must not shift haraqat + # onto the wrong letters; dropped region emits bare text letters + text = "كتب" + pred = "كَتَ" # missing final ب + out = project_haraqat(pred, text) + assert strip_diacritics(out) == text + assert out.startswith("كَتَ") + + +def test_project_haraqat_extra_prediction_letters() -> None: + text = "قل" + pred = "قَوْلٌ" # و inserted, haraqat on others + out = project_haraqat(pred, text) + assert strip_diacritics(out) == text + + +def test_diacritics_regex_covers_haraqat_and_shadda() -> None: + for ch in "ًٌٍَُِّْٰ": + assert DIACRITICS_RE.match(ch), ch