oddadmix/Nawah-Router-v3

1"""2Arabic zero-shot router with a real routing head.3 4Two earlier formulations, and why this one:5 6  fixed-slot head  - `num_labels = max_lanes` over one pooled vector. Scored exactly 1/n at every7                     lane count. The head is positional (slot i's logit is w_i . h) and the corpus8                     randomises lane order, so there is nothing to learn.9  pairwise         - score each (text, category) separately, argmax. Works (0.842 on unseen lane10                     sets) but costs n forward passes and each category is scored in isolation,11                     never against its competitors.12 13Here the categories live in the sequence and get their *own* pooled vectors, which a shared linear14head turns into one logit each; the softmax is over the categories present. The head is shared15across positions, so it scores category *content*, not slot index - which is what makes it16position-invariant and zero-shot.17 18Layout is text-first, categories-second, deliberately. The reference encoder is bidirectional so19order does not matter there; our base is causal, and this ordering is what lets every category20token attend to the whole text. Reversed, the categories would be encoded blind to the text.21 22  النص:23  {text}24 25  الفئات:26  - فئة أولى27  - فئة ثانية28 29Each category's character span is mapped to token indices via the tokenizer's offset mapping and30mean-pooled. Slots beyond a row's category count are masked to -inf before the loss.31"""32import json33import os34from collections import defaultdict35from pathlib import Path36 37import numpy as np38import torch39import torch.nn as nn40import torch.nn.functional as F41from torch.utils.data import Dataset42from transformers import AutoModel, AutoTokenizer, Trainer, TrainingArguments43 44os.environ.setdefault("CUDA_VISIBLE_DEVICES", "0")45 46BASE_MODEL = os.environ.get("BASE_MODEL", "oddadmix/50M-2048-Emhotob")47OUTPUT_DIR = os.environ.get("OUTPUT_DIR", "./Nawah-Router-head-50M")48DATA = Path(os.environ.get("DATA_DIR", "data"))49MAX_ROUTES = int(os.environ.get("MAX_ROUTES", 9))50MAX_LENGTH = int(os.environ.get("MAX_LENGTH", 320))51LR = float(os.environ.get("LR", 3e-4))52EPOCHS = float(os.environ.get("EPOCHS", 3))53BATCH_SIZE = int(os.environ.get("BATCH_SIZE", 32))54WARMUP = int(os.environ.get("WARMUP", 500))55SEED = 4256 57 58def build_text(row):59    """-> (full string, [(start, end) char span per category])."""60    head = f"النص:\n{row['text']}\n\nالفئات:\n"61    s = head62    spans = []63    for c in row["routes"]:64        s += "- "65        spans.append((len(s), len(s) + len(c)))66        s += c + "\n"67    return s, spans68 69 70class RouterModel(nn.Module):71    """Backbone + a shared linear scorer applied to each category's pooled span."""72 73    def __init__(self, base_model):74        super().__init__()75        self.backbone = AutoModel.from_pretrained(base_model, dtype=torch.float32)76        h = self.backbone.config.hidden_size77        self.score = nn.Sequential(nn.Linear(h, h), nn.GELU(), nn.Linear(h, 1))78        self.config = self.backbone.config79 80    def forward(self, input_ids, attention_mask, cat_pool, n_routes, labels=None):81        hs = self.backbone(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state82        # cat_pool: (B, K, T) row-normalised selector -> (B, K, H)83        cat_vecs = torch.bmm(cat_pool.to(hs.dtype), hs)84        logits = self.score(cat_vecs).squeeze(-1)                       # (B, K)85        ar = torch.arange(logits.size(1), device=logits.device)[None, :]86        logits = logits.masked_fill(ar >= n_routes[:, None], torch.finfo(logits.dtype).min)87        loss = F.cross_entropy(logits, labels) if labels is not None else None88        return {"loss": loss, "logits": logits}89 90 91def load(name):92    p = DATA / f"{name}.jsonl"93    return [json.loads(l) for l in open(p, encoding="utf-8")] if p.exists() else []94 95 96class RouterDataset(Dataset):97    def __init__(self, rows, tok, max_length):98        self.rows, self.tok, self.max_length = rows, tok, max_length99 100    def __len__(self):101        return len(self.rows)102 103    def __getitem__(self, i):104        r = self.rows[i]105        text, spans = build_text(r)106        enc = self.tok(text, return_offsets_mapping=True, add_special_tokens=False,107                       truncation=True, max_length=self.max_length)108        ids, offs = enc["input_ids"], enc["offset_mapping"]109        pool = torch.zeros(MAX_ROUTES, len(ids))110        for ci, (s, e) in enumerate(spans[:MAX_ROUTES]):111            idxs = [t for t, (a, b) in enumerate(offs) if a < e and b > s and a != b]112            if idxs:113                pool[ci, idxs] = 1.0 / len(idxs)114        return {"input_ids": torch.tensor(ids, dtype=torch.long), "cat_pool": pool,115                "n_routes": torch.tensor(min(len(r["routes"]), MAX_ROUTES), dtype=torch.long),116                "labels": torch.tensor(min(r["label"], MAX_ROUTES - 1), dtype=torch.long)}117 118 119class Collator:120    def __init__(self, pad_id):121        self.pad_id = pad_id122 123    def __call__(self, feats):124        T = max(f["input_ids"].size(0) for f in feats)125        ids, att, pools = [], [], []126        for f in feats:127            n = f["input_ids"].size(0); pad = T - n128            ids.append(torch.cat([f["input_ids"], torch.full((pad,), self.pad_id, dtype=torch.long)]))129            att.append(torch.cat([torch.ones(n, dtype=torch.long), torch.zeros(pad, dtype=torch.long)]))130            pools.append(F.pad(f["cat_pool"], (0, pad)))131        return {"input_ids": torch.stack(ids), "attention_mask": torch.stack(att),132                "cat_pool": torch.stack(pools),133                "n_routes": torch.stack([f["n_routes"] for f in feats]),134                "labels": torch.stack([f["labels"] for f in feats])}135 136 137def metrics_fn(p):138    return {"accuracy": float((np.asarray(p.predictions).argmax(-1) ==139                               np.asarray(p.label_ids)).mean())}140 141 142@torch.no_grad()143def report(model, tok, rows, name, batch=64):144    model.eval()145    ds = RouterDataset(rows, tok, MAX_LENGTH); coll = Collator(tok.pad_token_id)146    preds = []147    for i in range(0, len(rows), batch):148        b = coll([ds[j] for j in range(i, min(i + batch, len(rows)))])149        b = {k: v.to(next(model.parameters()).device) for k, v in b.items()}150        b.pop("labels")151        preds += model(**b)["logits"].argmax(-1).tolist()152    correct = [int(p == r["label"]) for p, r in zip(preds, rows)]153    acc = sum(correct) / len(rows)154    rand = sum(1 / len(r["routes"]) for r in rows) / len(rows)155    print(f"\n[{name}]  n={len(rows):,}  route accuracy {acc:.4f}   random {rand:.4f}")156    out = {"accuracy": acc, "random_baseline": rand, "n": len(rows)}157    for key in ("n_routes", "difficulty", "mode"):158        if key == "mode" and not all("mode" in r for r in rows):159            continue160        b = defaultdict(lambda: [0, 0])161        for c, r in zip(correct, rows):162            k = len(r["routes"]) if key == "n_routes" else r[key]163            b[k][1] += 1; b[k][0] += c164        print(f"    by {key:<10} " +165              "  ".join(f"{k}:{v[0]/v[1]:.3f}(n={v[1]})" for k, v in sorted(b.items(), key=str)))166        out[f"by_{key}"] = {str(k): {"acc": v[0]/v[1], "n": v[1]} for k, v in b.items()}167    return out168 169 170def main():171    tok = AutoTokenizer.from_pretrained(BASE_MODEL)172    if tok.pad_token_id is None:173        tok.pad_token = tok.eos_token174    model = RouterModel(BASE_MODEL)175    print(f"[*] {BASE_MODEL} | params {sum(p.numel() for p in model.parameters())/1e6:.2f}M "176          f"| hidden {model.config.hidden_size}")177 178    train = load("train")179    evals = {"unseen_lanes": load("eval_unseen_lanes"),180             "unseen_domain": load("eval_unseen_domain"),181             # v2 only: axes held out of training entirely. The strongest zero-shot test, so the182             # checkpoint is selected on it when it exists.183             "unseen_axis": load("eval_unseen_axis"),184             "hard": load("eval_hard")}185    evals = {k: v for k, v in evals.items() if v}186    print(f"[*] train {len(train):,} | " + " | ".join(f"{k} {len(v):,}" for k, v in evals.items()))187 188    args = TrainingArguments(189        output_dir=OUTPUT_DIR, num_train_epochs=EPOCHS,190        per_device_train_batch_size=BATCH_SIZE, per_device_eval_batch_size=64,191        learning_rate=LR, lr_scheduler_type="cosine", warmup_steps=WARMUP,192        max_grad_norm=1.0, bf16=True, logging_steps=200,193        eval_strategy="steps", eval_steps=1000, save_strategy="steps", save_steps=1000,194        save_total_limit=2, load_best_model_at_end=True,195        metric_for_best_model=("eval_unseen_axis_accuracy" if "unseen_axis" in evals196                               else "eval_unseen_domain_accuracy"), greater_is_better=True,197        report_to=[], seed=SEED, dataloader_num_workers=4, remove_unused_columns=False,198        label_names=["labels"])199 200    trainer = Trainer(model=model, args=args,201                      train_dataset=RouterDataset(train, tok, MAX_LENGTH),202                      eval_dataset={k: RouterDataset(v, tok, MAX_LENGTH)203                                    for k, v in evals.items() if v},204                      data_collator=Collator(tok.pad_token_id), compute_metrics=metrics_fn)205    trainer.train()206 207    Path(OUTPUT_DIR).mkdir(exist_ok=True)208    torch.save(model.state_dict(), Path(OUTPUT_DIR, "router_model.pt"))209    model.backbone.save_pretrained(OUTPUT_DIR); tok.save_pretrained(OUTPUT_DIR)210    results = {k: report(model, tok, v, k) for k, v in evals.items() if v}211    Path(OUTPUT_DIR, "train_metrics.json").write_text(json.dumps(212        {"results": results, "base_model": BASE_MODEL, "formulation": "routing_head",213         "log_history": trainer.state.log_history}, ensure_ascii=False, indent=2), encoding="utf-8")214    print(f"\n[+] done -> {OUTPUT_DIR}")215 216 217if __name__ == "__main__":218    main()219