From 92f436e3534d41a9635b49d47907aee62c95196d Mon Sep 17 00:00:00 2001 From: Ronald Tse Date: Fri, 9 Oct 2026 23:20:17 +0800 Subject: [PATCH] fix: BIL_BS env + padded-label loss masking (LSTM OOM + label smearing) --- train_bilstm_hf.py | 240 +++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 240 insertions(+) diff --git a/train_bilstm_hf.py b/train_bilstm_hf.py index e69de29..2662cb9 100644 --- a/train_bilstm_hf.py +++ b/train_bilstm_hf.py @@ -0,0 +1,240 @@ +"""WO32 — their literal architecture: BiLSTM haraqat tagger at full +silver scale (news-pure, 900K silver cap), the QCRI recipe class. + +Fadel/QCRI-style: char embeddings -> stacked BiLSTM -> per-position +haraqat classifier over the plane-combo inventory. No pretrained +backbone — from-scratch, exactly their class. The SOTA-stack layer +lives in run-033 (byt5-large plane) above it. + +Dual gate: SadeedDiac-25 windowed zero-skip + WikiNews-2024 multiref +(their 2.70 is the ID+OOD target class). + +Usage: + python /ckpt/run-034-bilstm/train_bilstm_hf.py +""" + +from __future__ import annotations + +import csv +import json +import os +import random +import re +import sys +import time +from pathlib import Path + +sys.path.insert(0, "/code") + +import torch +import torch.nn as nn +from torch.utils.data import DataLoader, Dataset + +import haraqat_planes as HP + +DATA = Path("/train_data") +RUN = os.environ.get("BIL_RUN", "run-034-bilstm") +CKPT = Path("/ckpt") / RUN +SILVER_CAP = int(os.environ.get("BIL_SILVER_CAP", "900000")) +EPOCHS = int(os.environ.get("BIL_EPOCHS", "6")) +BS = int(os.environ.get("BIL_BS", "64")) +EMB = 128 +HID = 512 +LAYERS = 3 +WINDOW = 1400 + +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 main() -> None: + CKPT.mkdir(parents=True, exist_ok=True) + rng = random.Random(42) + + # ---- corpus: news-pure + full silver (identical to run-033) ---- + lines: list[str] = [] + news = [l for l in (DATA / "news.txt").read_text(encoding="utf-8").splitlines() if l.strip()] + gold = [l for l in (DATA / "wikinews2014_gold.txt").read_text(encoding="utf-8").splitlines() if l.strip()] + units = news + gold * 4 + silver = 0 + articles = (DATA / "qcri-wiki.jsonl").read_text(encoding="utf-8").splitlines() + rng.shuffle(articles) + for aline in articles: + if silver >= SILVER_CAP: + break + if not aline.strip(): + continue + art = json.loads(aline).get("text", "") + for para in art.split("\n"): + para = para.strip() + if not para: + continue + for w in split_windows(para): + units.append(w) + silver += 1 + if silver >= SILVER_CAP: + break + print(f"[data] news={len(news)} gold={len(gold)} silver={silver}", flush=True) + + # defense: drop units whose SKELETON is empty (pure-mark windows + # would crash/produce nothing) — mirrors make_unit in the plane trainer + units = [u for u in (u.strip() for u in units) + if u and DIACRITICS_RE.sub("", u).encode("utf-8") != b""] + random.Random(42).shuffle(units) + classes, counts = HP.plane_classes(units) + class_to_id = {c: i for i, c in enumerate(classes)} + n_classes = len(classes) + print(f"[data] units={len(units)} classes={n_classes}", flush=True) + + # char vocab from corpus skeletons + vocab = {"": 0, "": 1} + for u in units: + for ch in HP.split_planes(u)[0]: + if ch not in vocab: + vocab[ch] = len(vocab) + print(f"[vocab] {len(vocab)}", flush=True) + + def encode(text: str): + skel, labels = HP.split_planes(text) + ids = [vocab.get(ch, 1) for ch in skel] + lab = [0] * len(ids) + for i, (ch, cluster) in enumerate(zip(skel, labels)): + lab[i] = class_to_id.get(HP.canon_combo(cluster), 0) + return ids, lab + + class DS(Dataset): + def __len__(self): + return len(units) + + def __getitem__(self, i): + return encode(units[i]) + + def collate(b): + mx = max(len(x[0]) for x in b) + ids = torch.zeros((len(b), mx), dtype=torch.long) + lab = torch.full((len(b), mx), -100, dtype=torch.long) # mask pads + for i, (x, y) in enumerate(b): + ids[i, :len(x)] = torch.tensor(x) + lab[i, :len(y)] = torch.tensor(y) + return ids, lab + + loader = DataLoader(DS(), batch_size=BS, shuffle=True, collate_fn=collate, + num_workers=2, drop_last=True) + total_steps = len(loader) * EPOCHS + + class BiLSTMTagger(nn.Module): + def __init__(self): + super().__init__() + self.emb = nn.Embedding(len(vocab), EMB, padding_idx=0) + self.lstm = nn.LSTM(EMB, HID, num_layers=LAYERS, bidirectional=True, + batch_first=True, dropout=0.3) + self.head = nn.Linear(HID * 2, n_classes) + + def forward(self, ids): + h, _ = self.lstm(self.emb(ids)) + return self.head(h) + + model = BiLSTMTagger().cuda().train() + n_params = sum(p.numel() for p in model.parameters()) + print(f"[model] params={n_params/1e6:.1f}M steps={total_steps}", flush=True) + + opt = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=0.01) + ce = nn.CrossEntropyLoss(ignore_index=-100) + + step = 0 + t0 = time.time() + for _ep in range(EPOCHS): + for ids, lab in loader: + ids, lab = ids.cuda(), lab.cuda() + with torch.autocast("cuda", torch.bfloat16): + logits = model(ids) + loss = ce(logits.reshape(-1, n_classes), lab.reshape(-1)) + loss.backward() + torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) + opt.step() + opt.zero_grad() + step += 1 + if step % 200 == 0: + print(f"[train] {step}/{total_steps} loss={float(loss):.3f} " + f"({step/(time.time()-t0):.1f} it/s)", flush=True) + + model.eval() + best = CKPT / "best" + best.mkdir(exist_ok=True) + torch.save({"model": model.state_dict(), "classes": classes, + "vocab": vocab}, best / "model.pt") + + # ---- gates ---- + import pandas as pd + import pyarrow.parquet as pq + + @torch.no_grad() + def tag(text: str) -> str: + skel, _ = HP.split_planes(text) + ids = [vocab.get(ch, 1) for ch in skel] or [0] + t = torch.tensor([ids], device="cuda") + with torch.autocast("cuda", torch.bfloat16): + logits = model(t) + preds = logits[0].argmax(-1).tolist() + return "".join(ch + (classes[preds[i]] if i < len(preds) else "") + for i, ch in enumerate(skel)) + + table = pq.read_table(DATA / "sadeeddiac-25.parquet") + s_inputs = [DIACRITICS_RE.sub("", t) for t in table.column("input").to_pylist()] + s_outputs = table.column("output").to_pylist() + all_windows, counts = [], [] + for text in s_inputs: + ws = split_windows(text) + counts.append(len(ws)) + all_windows.extend(ws) + preds = [tag(w) for w in all_windows] + k = 0 + paras = [] + for _t, c in zip(s_inputs, counts): + paras.append(" ".join(preds[k:k+c])); k += c + csv_p = Path("/tmp/bilstm_sadeed.csv") + pd.DataFrame({"gt": s_outputs, "pred": paras}).to_csv(csv_p, index=False, header=False) + (CKPT / "sadeed_preds.csv").write_text(csv_p.read_text(), encoding="utf-8") + + from sadeed_evaluator import ArabicDiacritizationEvaluator as E + + print(f"===== {RUN} SadeedDiac-25 (QCRI-target class) =====", flush=True) + E.report_errors_on_csv_file(str(csv_p), ground_truth_column_index=0, + predicted_column_index=1, has_header=False, + gt_missing_diacritic_is_error=False) + + import eval_wikinews_multiref as WN + + bench = [l for l in (DATA / "wikinews/WikiNews_2024_Multi_Ref.txt.diac") + .read_text(encoding="utf-8").splitlines() if l.strip()] + winputs = [DIACRITICS_RE.sub("", l) for l in bench] + ww, wc = [], [] + for text in winputs: + ws = split_windows(text); wc.append(len(ws)); ww.extend(ws) + wpreds = [tag(w) for w in ww] + k = 0; wparas = [] + for _t, c in zip(winputs, wc): + wparas.append(" ".join(wpreds[k:k+c])); k += c + wscore = WN.score(wparas, bench, skip_last=False) + print(json.dumps({"run": RUN, "surface": "wikinews-2024-multiref", **wscore}, + ensure_ascii=False), flush=True) + (CKPT / "wikinews_multiref.json").write_text( + json.dumps({"run": RUN, **wscore}, ensure_ascii=False, indent=1)) + (CKPT / "EVAL_DONE").write_text("ok\n") + + +if __name__ == "__main__": + main()