diff --git a/build_yallamorph_aux.py b/build_yallamorph_aux.py new file mode 100644 index 0000000..3d1c5a2 --- /dev/null +++ b/build_yallamorph_aux.py @@ -0,0 +1,93 @@ +"""YallaMorph/CamelMorph aux line builder for the r9 Arabic teacher. + +Generates feature-conditional vocalization pairs from paradigm +enumeration (deterministic dictionary resource — not LLM labels): + input = "MORPH: | " + target = "" +""" + +from __future__ import annotations + +import random +import re +from typing import Iterable, Iterator + +DIACRITICS_RE = re.compile("[ؐ-ًؚ-ٰٟۖ-ۜ۟-۪ۨ-ۭ]") + +FORM_MAX_BYTES = 200 +FEAT_KEY_ORDER = ("pos", "asp", "per", "gen", "num", "cas", "stt", "voz", "mod") +PERSONS = ("1", "2", "3") +GENDERS = ("m", "f") +NUMBERS = ("s", "d", "p") +CASES = ("nom", "acc", "gen") +STATES = ("d", "i") +VOICES = ("a", "p") +MOODS = ("i", "s", "j", "el", "eh") + +PERSON_TOKEN = {"1": "1st", "2": "2nd", "3": "3rd"} + + +def strip_diacritics(text: str) -> str: + return DIACRITICS_RE.sub("", text) + + +def feat_tokens(feat) -> str: + items = feat.items() if hasattr(feat, "items") else feat + by_key = dict(items) + ordered = [(k, by_key[k]) for k in FEAT_KEY_ORDER if k in by_key] + ordered += [(k, v) for k, v in by_key.items() if k not in FEAT_KEY_ORDER] + return " ".join(f"{k}:{v}" for k, v in ordered) + + +def aux_pair(feat, form: str) -> tuple[str, str] | None: + if not form: + return None + undiac = strip_diacritics(form) + if not undiac or len(form.encode("utf-8")) > FORM_MAX_BYTES: + return None + return f"MORPH: {feat_tokens(feat)} | {undiac}", form + + +def iter_verb_specs() -> Iterator[dict]: + for asp, per, gen, num, voz in ( + (asp, per, gen, num, voz) + for asp in ("pv", "iv", "c") + for per in (PERSONS if asp != "c" else ("2",)) + for gen in GENDERS + for num in NUMBERS + for voz in VOICES + ): + if asp == "iv": + for mod in MOODS: + yield {"pos": "v", "asp": asp, "per": PERSON_TOKEN[per], "gen": gen, "num": num, "voz": voz, "mod": mod} + else: + yield {"pos": "v", "asp": asp, "per": PERSON_TOKEN[per], "gen": gen, "num": num, "voz": voz} + + +def _iter_decl_specs(pos: str) -> Iterator[dict]: + for gen in GENDERS: + for num in NUMBERS: + for cas in CASES: + for stt in STATES: + yield {"pos": pos, "gen": gen, "num": num, "cas": cas, "stt": stt} + + +def iter_noun_specs() -> Iterator[dict]: + yield from _iter_decl_specs("n") + + +def iter_adj_specs() -> Iterator[dict]: + yield from _iter_decl_specs("adj") + + +def build_lines(rows: Iterable[tuple], cap: int, seed: int = 42) -> list[str]: + seen: set[tuple[str, str]] = set() + pairs: list[tuple[str, str]] = [] + for feat, form in rows: + p = aux_pair(feat, form) + if p is None or p in seen: + continue + seen.add(p) + pairs.append(p) + random.Random(seed).shuffle(pairs) + return [f"{src}\t{tgt}" for src, tgt in pairs[:cap]] diff --git a/gen_yallamorph_pairs.py b/gen_yallamorph_pairs.py new file mode 100644 index 0000000..ebb0dd0 --- /dev/null +++ b/gen_yallamorph_pairs.py @@ -0,0 +1,331 @@ +"""Generate the r9 MORPH aux corpus from CamelMorph MSA (CAMeL Tools). + +Pipeline: +1. Load the YallaMorph lemma sample (4,795 lemmas; xlsx sample files). +2. Enumerate feature specs per POS (baseword + clitic slots). +3. Generate diacritized forms with camel_tools Generator. +4. Validate against YallaMorph few-shot gold examples (prompt files). +5. Emit build_lines output (srctarget) + stats. + +Usage: + /tmp/camelenv/bin/python gen_yallamorph_pairs.py \ + --yalla-dir /tmp/yallamorph --out data/yallamorph-aux/lines.txt +""" + +from __future__ import annotations + +import argparse +import json +import random +import re +import sys +from collections import Counter +from concurrent.futures import ProcessPoolExecutor +from pathlib import Path + +from build_yallamorph_aux import ( + aux_pair, + build_lines, + iter_adj_specs, + iter_noun_specs, + iter_verb_specs, +) + +# YallaMorph prompt vocabulary -> calima-msa-r13 generation-DB feature codes +Y2C = { + "pos": {"verb": "verb", "noun": "noun", "adjective": "adj", "adj": "adj", + "v": "verb", "n": "noun"}, + "per": {"1st": "1", "2nd": "2", "3rd": "3", "1": "1", "2": "2", "3": "3"}, + "gen": {"m": "m", "f": "f", "Masculine": "m", "Feminine": "f"}, + "num": {"s": "s", "d": "d", "p": "p", "Singular": "s", "Dual": "d", "Plural": "p"}, + "asp": {"pv": "p", "iv": "i", "c": "c", "Perfective": "p", "Imperfective": "i", "Command": "c"}, + "mod": {"i": "i", "s": "s", "j": "j", "el": "e", "eh": "e", + "Indicative": "i", "Subjunctive": "s", "Jussive": "j", + "ENERGETIC_LIGHT": "e", "ENERGETIC_HEAVY": "e"}, + "voz": {"a": "a", "p": "p", "Active": "a", "Passive": "p"}, + "cas": {"nom": "n", "acc": "a", "gen": "g", + "Nominative": "n", "Accusative": "a", "Genitive": "g"}, + "stt": {"d": "d", "i": "i", "c": "c", "definite": "d", "indefinite": "i", "construct": "c"}, +} + +# clitic slots (YallaMorph name -> calima slot); values = calima-msa-r13 vocabulary +PROC_SLOTS = { + "proclitics_0": "prc0", + "proclitics_1": "prc1", + "proclitics_2": "prc2", + "proclitics_3": "prc3", +} +PROC_VALUES = { + "proclitics_0": ("0", "Al_det"), + "proclitics_1": ("0", "bi_prep", "li_prep", "ka_prep", "fiy_prep", "la_prep"), + "proclitics_2": ("0", "wa_conj", "fa_conj"), + "proclitics_3": ("0", ">a_ques"), +} +ENC_VALUES = ("0", "1s_poss", "2ms_poss", "2fs_poss", "3ms_poss", "3fs_poss", "1p_poss", "3mp_poss", "3fp_poss") + + +def calima_spec(spec: dict) -> dict: + out = {} + for k, v in spec.items(): + k = "vox" if k == "voz" else k + if k in Y2C: + if v in Y2C[k]: + out[k] = Y2C[k][v] + else: + out[k] = v + return out + + +# Camel Morph MSA v1.0 (LREC-COLING 2024 release, CC BY 4.0) — the +# resource YallaMorph was constructed from; usable directly by +# camel_tools MorphologyDB. Fetch: +# git clone https://github.com/CAMeL-Lab/camel_morph +DB_PATH = ("/tmp/camel_morph/official_releases/lrec-coling2024_release" + "/databases/camel-morph-msa/camel_morph_msa_v1.0.db") + + +def make_generator(db_path: str = DB_PATH): + from camel_tools.morphology.database import MorphologyDB + from camel_tools.morphology.generator import Generator + + return Generator(MorphologyDB(db_path, flags="g")) + + +def iter_noun_clitic_specs() -> list[dict]: + """Bounded clitic variants — the full cross product explodes (14 extra + per base) and would overrun the 300k line cap 20x over.""" + base = list(iter_noun_specs()) + list(iter_adj_specs()) + enc = ENC_VALUES[1:] + specs = [] + for i, s in enumerate(base): + specs.append(s) + specs.append({**s, "prc0": "Al_det"}) + specs.append({**s, "prc1": "bi_prep"}) + specs.append({**s, "prc0": "Al_det", "prc2": "wa_conj"}) + specs.append({**s, "enc0": enc[i % len(enc)]}) + return specs + + +def iter_verb_clitic_specs() -> list[dict]: + base = [s for s in iter_verb_specs()] + specs = list(base) + for s in base: + if s["asp"] == "iv": + for p1 in ("bi_prep", "la_prep", "sa_fut"): + specs.append({**s, "prc1": p1}) + for p2 in ("wa_conj", "fa_conj"): + specs.append({**s, "prc2": p2}) + return specs + + +def load_lemmas(yalla_dir: Path) -> dict[str, set[str]]: + import openpyxl + + pos_files = { + "v": ["Base Sample/Verb", "Clitics Sample/Verb"], + "n": ["Base Sample/Noun", "Clitics Sample/Noun"], + "adj": ["Base Sample/Adj", "Clitics Sample/Adj"], + } + lemmas: dict[str, set[str]] = {"v": set(), "n": set(), "adj": set()} + for pos, subdirs in pos_files.items(): + for sub in subdirs: + d = yalla_dir / "YallaMorph_Data" / sub + if not d.exists(): + continue + for x in sorted(d.glob("*.xlsx")): + ws = openpyxl.load_workbook(x, read_only=True).active + for row in ws.iter_rows(min_row=2, values_only=True): + if row and row[0]: + lemmas[pos].add(str(row[0]).strip()) + return lemmas + + +def parse_fewshot_gold(yalla_dir: Path) -> list[dict]: + """(input feature dict, set of gold forms) from the few-shot prompts.""" + gold = [] + prompts = yalla_dir / "Prompts" / "English_prompts" / "Few_shots" + for p in sorted(prompts.glob("*.txt")): + text = p.read_text(encoding="utf-8") + for m in re.finditer(r"Input:\n(.*?)\nOutput:\n(.*?)\n\n", text, re.S): + feat, out = {}, {} + for line in m.group(1).strip().splitlines(): + if ":" in line: + k, v = line.split(":", 1) + feat[k.strip()] = v.strip() + try: + out = json.loads(m.group(2).strip()) + except json.JSONDecodeError: + continue + if out.get("status") == "OK" and out.get("arabic_forms"): + gold.append({"file": p.name, "feat": feat, "forms": set(out["arabic_forms"])}) + return gold + + +_GEN = None + + +def _get_generator(): + global _GEN + if _GEN is None: + _GEN = make_generator() + return _GEN + + +# Underspecified generation is the correct interface: the DB's merged +# analyses return the complete paradigm with per/gen/num/mod filled by +# the suffix tables (e.g. 1st person carries gen 'u', dual 2nd 'u'), +# which fully-specified requests silently reject. One call per variant +# replaces the full spec cross-product and yields exactly the cells +# the resource knows. +VERB_VARIANTS = [ + {}, + {"prc1": "bi_prep"}, + {"prc1": "li_prep"}, + {"prc1": "sa_fut"}, + {"prc2": "wa_conj"}, + {"prc2": "fa_conj"}, + {"prc3": ">a_ques"}, +] +NOUN_VARIANTS = [ + {}, + {"prc0": "Al_det"}, + {"prc1": "bi_prep"}, + {"prc0": "Al_det", "prc2": "wa_conj"}, + {"enc0": "1s_poss"}, + {"enc0": "3ms_poss"}, + {"enc0": "1p_poss"}, +] +VERB_KEYS = ("pos", "asp", "per", "gen", "num", "vox", "mod", "prc0", "prc1", "prc2", "prc3", "enc0") +NOUN_KEYS = ("pos", "gen", "num", "cas", "stt", "prc0", "prc1", "prc2", "prc3", "enc0") + + +def _worker(args: tuple[str, str]) -> list[tuple[dict, str]]: + lex, pos = args + gen = _get_generator() + variants = VERB_VARIANTS if pos == "v" else NOUN_VARIANTS + keys = VERB_KEYS if pos == "v" else NOUN_KEYS + db_pos = {"v": "verb", "adj": "adj", "n": "noun"}[pos] + rows = [] + for extra in variants: + spec = {"pos": db_pos, **extra} + try: + analyses = gen.generate(lex, spec) + except Exception: + continue + for a in analyses: + form = a.get("diac") + if not form: + continue + feat = {k: a[k] for k in keys if a.get(k) not in (None, "", "na", "0")} + rows.append((feat, form)) + return rows + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--yalla-dir", type=Path, default=Path("/tmp/yallamorph")) + ap.add_argument("--out", type=Path, default=Path("data/yallamorph-aux/lines.txt")) + ap.add_argument("--cap", type=int, default=300_000) + ap.add_argument("--workers", type=int, default=8) + ap.add_argument("--validate-only", action="store_true") + args = ap.parse_args() + + lemmas = load_lemmas(args.yalla_dir) + print({k: len(v) for k, v in lemmas.items()}) + + gen = make_generator() + in_lex = {pos: {lex for lex in lm if lex in gen._db.lemma_hash} for pos, lm in lemmas.items()} + print("yallamorph-sample in-lexicon:", {k: len(v) for k, v in in_lex.items()}) + + TARGETS = {"v": 4000, "n": 6000, "adj": 3000} + db_by_pos = {"verb": "v", "noun": "n", "adj": "adj"} + inventory: dict[str, list[str]] = {p: [] for p in db_by_pos.values()} + for lex, stems in gen._db.lemma_hash.items(): + p = db_by_pos.get(stems[0].get("pos", "")) + if p: + inventory[p].append(lex) + + selected: dict[str, set[str]] = {} + for pos, target in TARGETS.items(): + chosen = set(in_lex[pos]) + rest = sorted(set(inventory[pos]) - chosen) + chosen.update(random.Random(42).sample(rest, max(0, target - len(chosen)))) + selected[pos] = chosen + print("selected lemmas:", {k: len(v) for k, v in selected.items()}) + + gold = parse_fewshot_gold(args.yalla_dir) + print(f"[validate] {len(gold)} gold few-shot examples") + n_ok = n_tot = 0 + for g in gold: + feat = g["feat"] + spec = { + "pos": Y2C["pos"].get(feat.get("POS", ""), ""), + "gen": Y2C["gen"].get(feat.get("GENDER", ""), ""), + "num": Y2C["num"].get(feat.get("NUMBER", ""), ""), + } + if feat.get("PERSON"): + spec["per"] = Y2C["per"][feat["PERSON"]] + if feat.get("ASPECT"): + spec["asp"] = Y2C["asp"][feat["ASPECT"]] + if feat.get("MOD"): + spec["mod"] = Y2C["mod"][feat["MOD"]] + if feat.get("VOICE"): + spec["vox"] = Y2C["voz"][feat["VOICE"]] + if feat.get("CASE"): + spec["cas"] = Y2C["cas"][feat["CASE"]] + if feat.get("STATE"): + spec["stt"] = Y2C["stt"][feat["STATE"]] + for pname in sorted(PROC_SLOTS): + v = feat.get(pname, "0") + if v != "0": + key = {"Al_det": "prc0", "bi_prep": "prc1", "li_prep": "prc1", + "wa_conj": "prc2", "fa_conj": "prc2", ">a_ques": "prc3"}.get(v) + if key: + spec[key] = v + for pname in sorted(feat): + if pname.startswith("enclitics_0") and feat[pname] != "0": + spec["enc0"] = feat[pname] + spec = {k: v for k, v in spec.items() if v} + try: + forms = {a.get("diac") for a in gen.generate(feat["LEMMA"], spec)} + except Exception as e: + print(f"[validate] ERROR {g['file']}: {spec} -> {e}") + continue + hit = bool(forms & g["forms"]) + n_ok += hit + n_tot += 1 + if not hit: + print(f"[miss] {g['file']}: {spec} -> ours={sorted(f for f in forms if f)[:3]} gold={sorted(g['forms'])}") + print(f"[validate] exact-form match {n_ok}/{n_tot}") + if args.validate_only: + return + if n_tot == 0 or n_ok / n_tot < 0.5: + print("[abort] validation below 50% — fix feature mapping first") + sys.exit(1) + + jobs = [(lex, pos) for pos, lm in selected.items() for lex in sorted(lm)] + print(f"[gen] {len(jobs)} lemmas (underspecified paradigm generation)") + + rows: list[tuple[dict, str]] = [] + with ProcessPoolExecutor(max_workers=args.workers) as ex: + for i, part in enumerate(ex.map(_worker, jobs, chunksize=8)): + rows.extend(part) + if (i + 1) % 500 == 0: + print(f"[gen] {i + 1}/{len(jobs)} lemmas -> {len(rows)} pairs", flush=True) + print(f"[gen] {len(rows)} raw (feat, form) pairs") + + verb_rows = [r for r in rows if r[0].get("pos") == "verb"] + nom_rows = [r for r in rows if r[0].get("pos") != "verb"] + verb_cap = int(args.cap * 0.6) + nom_cap = args.cap - verb_cap + lines = build_lines(verb_rows, cap=verb_cap) + build_lines(nom_rows, cap=nom_cap) + args.out.parent.mkdir(parents=True, exist_ok=True) + args.out.write_text("\n".join(lines) + "\n", encoding="utf-8") + pos_counts = Counter() + for spec, _ in rows: + pos_counts[spec.get("pos", "?")] += 1 + print(f"[out] {args.out} lines={len(lines)} raw-by-pos={dict(pos_counts)}") + + +if __name__ == "__main__": + main() diff --git a/probe_ride_direction.py b/probe_ride_direction.py new file mode 100644 index 0000000..20cb9cf --- /dev/null +++ b/probe_ride_direction.py @@ -0,0 +1,118 @@ +"""RIDE direction-transferability probe (TODO.sota-2026/05). + +Question: is the r7-over-r6 SFT residual direction (per encoder layer) +domain-general, or news-idiosyncratic? RIDE-style extrapolation +(regressing a student toward h_teacher + λ(h_teacher − h_base)) is only +sound if the direction transfers. + +Method: mean-pool per-layer encoder hidden states of BOTH teachers on +two text domains (classical SadeedDiac windows vs news lines), form +d_domain = mean_r7 − mean_r6, then cos(d_classical, d_news) per layer. + +Kill rule: max-layer cosine < 0.5 ⇒ direction domain-idiosyncratic ⇒ +close the arm without any training compute. + +Usage: + modal run probe_ride_direction.py +""" + +from __future__ import annotations + +import json +import re +from pathlib import Path + +import modal + +checkpoints_volume = modal.Volume.from_name("rababa-checkpoints", create_if_missing=True) +datasets_volume = modal.Volume.from_name("rababa-datasets", create_if_missing=True) + +BASE_RUN = "rababa_arabic_byt5/run-006-morph" +TEACHER_RUN = "rababa_arabic_byt5/run-007-news" +N_UNITS = 200 +KILL_COSINE = 0.5 + +DIACRITICS_RE = re.compile("[ؐ-ًؚ-ٰٟۖ-ۜ۟-۪ۨ-ۭ]") + +image = ( + modal.Image.debian_slim(python_version="3.11") + .pip_install("torch==2.5.1", "transformers==4.46.3", "pyarrow", "tqdm") + .add_local_dir("data/sadeed-diac-25", "/opt/rababa/data/sadeed-diac-25", copy=True) + .add_local_file("ride_probe_math.py", "/opt/rababa/ride_probe_math.py", copy=True) + .workdir("/opt/rababa") +) + +app = modal.App("rababa-ride-probe", image=image) + + +@app.function(gpu="A10G", timeout=2 * 60 * 60, volumes={"/checkpoints": checkpoints_volume, "/datasets": datasets_volume}) +def probe() -> dict: + import pyarrow.parquet as pq + import torch + from transformers import AutoModelForSeq2SeqLM, AutoTokenizer + + from ride_probe_math import cosine_rows, residual_directions + + checkpoints_volume.reload() + datasets_volume.reload() + + table = pq.read_table("data/sadeed-diac-25/train.parquet") + classical = [DIACRITICS_RE.sub("", t) for t in table.column("input").to_pylist()[:N_UNITS]] + news = [l.strip() for l in (Path("/datasets/arabic-news-r5/news.txt")).read_text(encoding="utf-8").splitlines() if l.strip()][:N_UNITS] + + def layer_means(model_dir: str, texts: list[str]) -> list[list[float]]: + tok = AutoTokenizer.from_pretrained(model_dir) + model = AutoModelForSeq2SeqLM.from_pretrained(model_dir).cuda().eval() + sums, counts = None, 0 + with torch.no_grad(): + for i in range(0, len(texts), 8): + batch = [t[:1400] for t in texts[i : i + 8]] + enc = tok(batch, return_tensors="pt", padding=True, truncation=True, max_length=1600).to("cuda") + mask = enc["attention_mask"].unsqueeze(-1).float() + out = model.encoder(**enc, output_hidden_states=True) + hs = out.hidden_states # tuple: (n_layers+1, B, T, D) + if sums is None: + sums = [torch.zeros(h.shape[-1], device="cuda") for h in hs] + for li, h in enumerate(hs): + sums[li] += (h * mask).sum(dim=(0, 1)) + counts += int(enc["attention_mask"].sum().item()) + del model + torch.cuda.empty_cache() + return [(s / counts).tolist() for s in sums] + + base_dir = str(Path("/checkpoints") / BASE_RUN / "best") + teach_dir = str(Path("/checkpoints") / TEACHER_RUN / "best") + + results = {} + for domain, texts in (("classical", classical), ("news", news)): + b = layer_means(base_dir, texts) + t = layer_means(teach_dir, texts) + results[domain] = residual_directions(b, t) + + cos = cosine_rows(results["classical"], results["news"]) + table_rows = [(i, c) for i, c in enumerate(cos)] + best = max(cos) + verdict = "TRANSFER" if best >= KILL_COSINE else "KILL" + print("layer cosine(classical-dir, news-dir):") + for i, c in table_rows: + print(f" L{i:02d} {c:+.4f}") + print(f"max = {best:+.4f} -> {verdict} (kill < {KILL_COSINE})") + + out = { + "base": BASE_RUN, + "teacher": TEACHER_RUN, + "n_units": N_UNITS, + "cosine_by_layer": cos, + "max_cosine": best, + "kill_cosine": KILL_COSINE, + "verdict": verdict, + } + dest = Path("/checkpoints") / "ride_probe_r6_r7.json" + dest.write_text(json.dumps(out, indent=2), encoding="utf-8") + checkpoints_volume.commit() + return out + + +@app.local_entrypoint() +def main(): + print(probe.remote()) diff --git a/ride_probe_math.py b/ride_probe_math.py new file mode 100644 index 0000000..dfb27f8 --- /dev/null +++ b/ride_probe_math.py @@ -0,0 +1,32 @@ +"""Pure math for the RIDE direction-transferability probe. + +Input shapes are per-layer mean-pooled hidden states: + base[layer] = mean-pooled hidden vector of the base checkpoint + teacher[layer] = same for the improved checkpoint +The probe asks whether the residual direction + d = teacher - base +computed on one text domain agrees with the direction computed on +another — i.e. whether the SFT delta generalizes or is +domain-idiosyncratic. Kill rule (TODO.sota-2026/05): max-layer cosine +< 0.5 closes the arm. +""" + +from __future__ import annotations + +import math + + +def residual_directions(base: list[list[float]], teacher: list[list[float]]) -> list[list[float]]: + return [[t - b for b, t in zip(brow, trow)] for brow, trow in zip(base, teacher)] + + +def cosine_rows(a: list[list[float]], b: list[list[float]]) -> list[float]: + out = [] + for arow, brow in zip(a, b): + na = math.sqrt(sum(x * x for x in arow)) + nb = math.sqrt(sum(x * x for x in brow)) + if na == 0.0 or nb == 0.0: + out.append(0.0) + else: + out.append(sum(x * y for x, y in zip(arow, brow)) / (na * nb)) + return out diff --git a/test_ride_probe_math.py b/test_ride_probe_math.py new file mode 100644 index 0000000..b960180 --- /dev/null +++ b/test_ride_probe_math.py @@ -0,0 +1,46 @@ +"""Tests for the RIDE direction-transferability probe math.""" + +import math +import unittest +from pathlib import Path +import sys + +sys.path.insert(0, str(Path(__file__).resolve().parent)) + +from ride_probe_math import cosine_rows, residual_directions + + +class TestResidualDirections(unittest.TestCase): + def test_subtraction_per_layer(self): + base = [[1.0, 2.0], [0.0, 1.0]] + teach = [[3.0, 5.0], [2.0, 1.0]] + d = residual_directions(base, teach) + self.assertEqual(d, [[2.0, 3.0], [2.0, 0.0]]) + + def test_layer_count_preserved(self): + base = [[0.0] * 4 for _ in range(5)] + teach = [[1.0] * 4 for _ in range(5)] + self.assertEqual(len(residual_directions(base, teach)), 5) + + +class TestCosineRows(unittest.TestCase): + def test_identical_is_one(self): + self.assertAlmostEqual(cosine_rows([[1.0, 0.0]], [[2.0, 0.0]])[0], 1.0) + + def test_orthogonal_is_zero(self): + self.assertAlmostEqual(cosine_rows([[1.0, 0.0]], [[0.0, 3.0]])[0], 0.0) + + def test_opposite_is_minus_one(self): + self.assertAlmostEqual(cosine_rows([[1.0, 1.0]], [[-2.0, -2.0]])[0], -1.0) + + def test_zero_row_is_zero(self): + self.assertEqual(cosine_rows([[0.0, 0.0]], [[1.0, 1.0]])[0], 0.0) + + def test_rows_independent(self): + c = cosine_rows([[1.0, 0.0], [0.0, 1.0]], [[0.5, 0.5], [1.0, 0.0]]) + self.assertAlmostEqual(c[0], 1 / math.sqrt(2)) + self.assertAlmostEqual(c[1], 0.0) + + +if __name__ == "__main__": + unittest.main() diff --git a/test_yallamorph_aux.py b/test_yallamorph_aux.py new file mode 100644 index 0000000..60810d3 --- /dev/null +++ b/test_yallamorph_aux.py @@ -0,0 +1,96 @@ +"""Tests for the YallaMorph/CamelMorph aux line builder (r9 teacher). + +The aux stream teaches feature-conditional vocalization: + input = "MORPH: | " + target = "" +""" + +import sys +import unittest +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parent)) + +from build_yallamorph_aux import ( + aux_pair, + build_lines, + feat_tokens, + iter_adj_specs, + iter_noun_specs, + iter_verb_specs, + strip_diacritics, +) + + +class TestStrip(unittest.TestCase): + def test_strips_haraqat_keeps_base(self): + self.assertEqual(strip_diacritics("أَلْوَيْتُ"), "ألويت") + + def test_plain_untouched(self): + self.assertEqual(strip_diacritics("ألت"), "ألت") + + def test_empty_after_strip(self): + self.assertEqual(strip_diacritics("ًٌٍَُِّْ"), "") + + +class TestFeatTokens(unittest.TestCase): + def test_deterministic_order(self): + a = feat_tokens({"voz": "a", "asp": "pv", "gen": "m"}) + b = feat_tokens({"gen": "m", "asp": "pv", "voz": "a"}) + self.assertEqual(a, b) + self.assertEqual(a, "asp:pv gen:m voz:a") + + def test_ascii_only(self): + toks = feat_tokens({"asp": "pv", "per": "1", "num": "s"}) + self.assertTrue(all(ord(c) < 128 for c in toks)) + + +class TestAuxPair(unittest.TestCase): + FORM = "أَلْوَيْتُ" + + def test_prefix_and_undiacritized_input(self): + src, tgt = aux_pair({"asp": "pv", "per": "1", "gen": "m", "num": "s", "voz": "a"}, self.FORM) + self.assertTrue(src.startswith("MORPH: ")) + self.assertIn(" | ", src) + self.assertTrue(src.endswith("| ألويت")) + self.assertNotIn("َ", src.split("| ", 1)[1]) + self.assertEqual(tgt, self.FORM) + + def test_rejects_empty_form(self): + self.assertIsNone(aux_pair({"asp": "pv"}, "ًٌ")) + + def test_rejects_oversize_form(self): + self.assertIsNone(aux_pair({"asp": "pv"}, "ا" * 300)) + + +class TestSpecs(unittest.TestCase): + def test_verb_specs_shape(self): + specs = list(iter_verb_specs()) + self.assertGreater(len(specs), 0) + keys = {k for s in specs for k in s} + self.assertTrue(keys <= {"pos", "asp", "per", "gen", "num", "voz", "mod"}) + self.assertTrue(all(s.get("asp") in ("pv", "iv", "c") for s in specs)) + self.assertTrue(all("gen" in s and "num" in s and "per" in s for s in specs)) + + def test_noun_specs_shape(self): + for fn in (iter_noun_specs, iter_adj_specs): + specs = list(fn()) + self.assertGreater(len(specs), 0) + keys = {k for s in specs for k in s} + self.assertTrue(keys <= {"pos", "gen", "num", "cas", "stt"}) + self.assertTrue(all({"gen", "num", "cas", "stt"} <= set(s) for s in specs)) + + +class TestBuildLines(unittest.TestCase): + def test_dedupe_and_cap(self): + rows = [({"asp": "pv"}, "فَعَلَ")] * 3 + [({"asp": "iv"}, "يَفْعَلُ")] + lines = build_lines(rows, cap=10) + self.assertEqual(len(lines), 2) + + def test_cap_applies(self): + rows = [({"pos": "v", "asp": "pv", "per": str(i)}, "فَعَلَ") for i in range(100)] + self.assertEqual(len(build_lines(rows, cap=7)), 7) + + +if __name__ == "__main__": + unittest.main() diff --git a/train_arabic_r9_yallamorph.py b/train_arabic_r9_yallamorph.py new file mode 100644 index 0000000..ea15eee --- /dev/null +++ b/train_arabic_r9_yallamorph.py @@ -0,0 +1,350 @@ +"""Arabic r9 — YallaMorph/CamelMorph aux stream (r8-generation teacher). + +r7 (news mix, 2.2864 ID / 17.38 OOD) is the canonical teacher; every +frontier move has been data-side (r5→r6→r7). r9 adds feature-complete +morphological paradigm coverage on top of r7's mix: a `MORPH: `-prefixed +aux stream built from CamelMorph MSA paradigm enumeration (build_ +yallamorph_aux.py; validated against YallaMorph gold samples, +arXiv 2609.10153). Deterministic dictionary resource — no LLM labels. + +Stream M: input = "MORPH: | ", +target = . Sampled to ~25% of the mix +(r6-proven aux dose). Inference contract unchanged: no prefix at +inference = plain diacritization. + +Init: run-007-news. Batch 2/accum 15 (r5/r6/r7-proven at 1400B). + +Gates (canonical replacement requires BOTH): +- SadeedDiac-25 windowed zero-skip DER within +0.1 of r7 (<= 2.389); +- WikiNews-2024 multi-ref improves over 17.3794/11.8273 + (eval_wikinews_multiref.py --model-dir this run's best). + +Usage: + modal run --detach train_arabic_r9_yallamorph.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_byt5/run-009-yallamorph" +DEFAULT_INIT = "rababa_arabic_byt5/run-007-news" +UNIT_BYTES = 1400 +NEWS_UPSAMPLE = 3 +GOLD_UPSAMPLE = 4 +AUX_SHARE = 0.25 +N_VAL = 2_000 + +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_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-r9", image=image) + + +def make_pair(unit: str) -> tuple[str, str] | None: + src = DIACRITICS_RE.sub("", unit) + if not src: + return None + if len(src.encode("utf-8")) > 1450 or len(unit.encode("utf-8")) > 1450: + return None + return src, unit + + +def load_morph_pairs() -> list[tuple[str, str]]: + """lines.txt rows are 'srctarget' from build_yallamorph_aux.py.""" + path = Path("/datasets/yallamorph-aux/lines.txt") + pairs = [] + for line in path.read_text(encoding="utf-8").splitlines(): + if "\t" not in line: + continue + src, tgt = line.split("\t", 1) + if src and tgt: + pairs.append((src, tgt)) + return pairs + + +@app.function( + gpu="A100-80GB", + timeout=24 * 60 * 60, + volumes={"/datasets": datasets_volume, "/checkpoints": checkpoints_volume}, +) +def train() -> dict: + import torch + from torch.utils.data import Dataset + from transformers import ( + AutoModelForSeq2SeqLM, + AutoTokenizer, + DataCollatorForSeq2Seq, + Seq2SeqTrainer, + Seq2SeqTrainingArguments, + TrainerCallback, + ) + + datasets_volume.reload() + checkpoints_volume.reload() + + done_marker = Path("/checkpoints") / RUN / "EVAL_DONE" + if done_marker.exists(): + return {"run": RUN, "status": "already-done"} + + news_dir = Path("/datasets/arabic-news-r5") + if not (news_dir / "DONE").exists(): + return {"run": RUN, "status": "waiting-for-label_arabic_news"} + aux_dir = Path("/datasets/yallamorph-aux") + if not (aux_dir / "DONE").exists(): + return {"run": RUN, "status": "waiting-for-yallamorph-aux"} + + init_run = DEFAULT_INIT + if not (Path("/checkpoints") / init_run / "EVAL_DONE").exists(): + return {"run": RUN, "status": f"init-not-ready: {init_run}"} + print(f"[init] {init_run}", flush=True) + + print("[data] loading cached r5 paragraph units...", flush=True) + cache = Path("/datasets/r5-units") + domain = [l for l in (cache / "domain.txt").read_text(encoding="utf-8").splitlines() if l.strip()] + replay = [l for l in (cache / "replay.txt").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()] + print(f"[data] r5-units={len(domain) + len(replay)} news={len(news)} gold2014={len(gold)}", flush=True) + + pairs: list[tuple[str, str]] = [] + for unit in domain + replay: + p = make_pair(unit) + if p: + pairs.append(p) + n_anchor = len(pairs) + for unit in news * NEWS_UPSAMPLE + gold * GOLD_UPSAMPLE: + p = make_pair(unit) + if p: + pairs.append(p) + n_plain = len(pairs) + + morph = load_morph_pairs() + target_aux = int(n_plain * AUX_SHARE / (1 - AUX_SHARE)) + if len(morph) > target_aux: + morph = random.Random(43).sample(morph, target_aux) + print(f"[data] anchor={n_anchor} news-mix={n_plain - n_anchor} " + f"morph-avail={len(load_morph_pairs())} morph-used={len(morph)} " + f"aux-share={len(morph) / (n_plain + len(morph)):.2%}", flush=True) + pairs.extend(morph) + random.Random(42).shuffle(pairs) + + combined = [l for l in Path("/datasets/arabic-combined/train.txt").read_text(encoding="utf-8").splitlines() if l.strip()] + random.Random(42).shuffle(combined) + val_pairs = [p for p in (make_pair(l) for l in combined[:N_VAL]) if p][:200] + + init = str(Path("/checkpoints") / init_run / "best") + tokenizer = AutoTokenizer.from_pretrained(init) + model = AutoModelForSeq2SeqLM.from_pretrained(init) + + class UnitDataset(Dataset): + def __init__(self, rows: list[tuple[str, str]]) -> None: + self.rows = rows + + def __len__(self) -> int: + return len(self.rows) + + def __getitem__(self, idx: int) -> dict: + src, tgt = self.rows[idx] + inputs = tokenizer(src, truncation=True, max_length=1600) + labels = tokenizer(tgt, truncation=True, max_length=1600) + inputs["labels"] = labels["input_ids"] + return inputs + + class VolumeCommitCallback(TrainerCallback): + def on_save(self, args, state, control, **kwargs): + try: + checkpoints_volume.commit() + print(f"[volume] committed at step {state.global_step}", flush=True) + except Exception as e: + print(f"[volume] commit failed at step {state.global_step}: {e}", flush=True) + + args = Seq2SeqTrainingArguments( + output_dir="/checkpoints/" + RUN, + num_train_epochs=1, + per_device_train_batch_size=2, + gradient_accumulation_steps=15, + per_device_eval_batch_size=1, + bf16=True, + learning_rate=2e-5, + lr_scheduler_type="cosine", + warmup_steps=200, + weight_decay=0.01, + max_grad_norm=1.0, + label_smoothing_factor=0.1, + seed=42, + save_strategy="steps", + save_steps=300, + eval_strategy="epoch", + save_total_limit=1, + logging_steps=100, + report_to=[], + predict_with_generate=False, + dataloader_num_workers=4, + ) + trainer = Seq2SeqTrainer( + model=model, + args=args, + train_dataset=UnitDataset(pairs), + eval_dataset=UnitDataset(val_pairs), + data_collator=DataCollatorForSeq2Seq(tokenizer=tokenizer, model=model, label_pad_token_id=-100), + callbacks=[VolumeCommitCallback()], + ) + import glob + + latest = sorted(glob.glob(f"/checkpoints/{RUN}/checkpoint-*"), key=lambda p: int(p.rsplit("-", 1)[1])) + resume = latest[-1] if latest else None + print(f"[resume] {resume}", flush=True) + trainer.train(resume_from_checkpoint=resume) + + best = Path("/checkpoints") / RUN / "best" + best.mkdir(parents=True, exist_ok=True) + trainer.save_model(str(best)) + tokenizer.save_pretrained(str(best)) + checkpoints_volume.commit() + + # ---- ID gate: windowed zero-skip SadeedDiac (r7 harness verbatim) ---- + import pandas as pd + import pyarrow.parquet as pq + from difflib import SequenceMatcher + + 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() + + 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 project_haraqat(pred: str, text: str) -> str: + 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) + + device = next(trainer.model.parameters()).device + trainer.model.eval() + + 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) + + import json as _json + prog = Path("/checkpoints") / RUN / "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) + + missing = [i for i in range(len(all_windows)) if i not in saved] + n_new = 0 + with torch.no_grad(), prog.open("a", encoding="utf-8") as prog_out: + for bi in range(0, len(missing), 8): + idxs = missing[bi : bi + 8] + batch = [all_windows[i] for i in idxs] + enc = tokenizer( + batch, return_tensors="pt", padding=True, truncation=True, max_length=1600 + ).to(device) + with torch.autocast("cuda", torch.bfloat16): + gen = trainer.model.generate(**enc, max_new_tokens=3200, num_beams=1) + batch_preds = tokenizer.batch_decode(gen, skip_special_tokens=True) + for i, pred in zip(idxs, batch_preds): + 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() + preds = [saved[i] for i in range(len(all_windows))] + + 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_r9_windowed.csv") + pd.DataFrame({"gt": outputs, "pred": paragraphs}).to_csv(csv_path, index=False, header=False) + (Path("/checkpoints") / RUN / "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===== r9 yallamorph, windowed zero-skip (ID gate; canonical iff <= r7 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, "init": init_run} + + +@app.local_entrypoint() +def main(): + handle = train.spawn() + print(f"spawned {handle.object_id}; completion = EVAL_DONE marker at " + f"rababa_checkpoints:rababa_arabic_byt5/run-009-yallamorph/EVAL_DONE", flush=True)