From 2adad85ce93128e0d5a4f73ed6559b2b0305ff6c Mon Sep 17 00:00:00 2001
From: Ronald Tse <ronald.tse@ribose.com>
Date: Thu, 8 Oct 2026 20:26:44 +0800
Subject: [PATCH] =?UTF-8?q?feat:=20byt5-large=20dedicated=20ID=20arm=20?=
 =?UTF-8?q?=E2=80=94=20R8=5FBACKBONE/R8=5FCORPUS=3Dr5/R8=5FSCRATCH=20knobs?=
 =?UTF-8?q?=20(WO27)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 train_arabic_r8_hf.py | 38 ++++++++++++++++++++++++++++++++++++--
 1 file changed, 36 insertions(+), 2 deletions(-)

diff --git a/train_arabic_r8_hf.py b/train_arabic_r8_hf.py
index 8ecee2b..f707e85 100644
--- a/train_arabic_r8_hf.py
+++ b/train_arabic_r8_hf.py
@@ -40,6 +40,8 @@
 TEACHER = Path("/ckpt/r7-best")
 import os
 MIX = os.environ.get("R8_MIX", "arwiki")  # arwiki | qcri | both
+BACKBONE = os.environ.get("R8_BACKBONE", "google/byt5-base")
+CORPUS = os.environ.get("R8_CORPUS", "tashkeela")  # tashkeela | r5
 RUN = os.environ.get("R8_RUN", f"run-024-arabic-r8-{MIX}")
 CKPT = Path(f"/ckpt/{RUN}")
 WINDOW = 600
@@ -101,6 +103,33 @@ def main() -> None:
     rng = random.Random(42)
 
     gold: list[tuple[str, str]] = []
+    if CORPUS == "r5":
+        # the r7 lineage corpus (r5 units + news), no silver
+        for name in ("domain.txt", "replay.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 * 2:
+                    break
+            if len(gold) >= GOLD_CAP * 2:
+                break
+        for line in (DATA / "news.txt").read_text(encoding="utf-8").splitlines()[:120_000]:
+            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))
+        print(f"[data] r5-corpus gold={len(gold)}", flush=True)
     gold_files = ("tashkeela-full/train-001.txt", "tashkeela-full/train-002.txt",
                   "tashkeela-full/train-003.txt", "tashkeela-full/val-001.txt")
     if os.environ.get("R8_GOLD", "tashkeela") == "none":
@@ -173,7 +202,7 @@ def add_silver(tgt: str) -> None:
     rng.shuffle(units)
     rng.shuffle(gold)
 
-    tok = AutoTokenizer.from_pretrained(TEACHER)
+    tok = AutoTokenizer.from_pretrained(BACKBONE if os.environ.get("R8_SCRATCH") == "1" else TEACHER)
 
     class DS(Dataset):
         def __len__(self):
@@ -203,7 +232,12 @@ def collate(b):
     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)
+    if os.environ.get("R8_SCRATCH") == "1":
+        init_from = BACKBONE
+    else:
+        init_from = TEACHER
+    print(f"[init] from {init_from}", flush=True)
+    model = AutoModelForSeq2SeqLM.from_pretrained(init_from, torch_dtype=torch.float32)
     model = model.cuda().train()
     model.gradient_checkpointing_enable()
     model.config.use_cache = False