Nawah-Router-v3: 52M, 0.9305 unseen lanes / 0.6976 unseen domains

Trained on the 233,720-row arabic-prompt-routing corpus with 12 routing axes.
Adds an unseen_axis eval: two routing dimensions absent from training entirely.

README.md

@@ -0,0 +1,75 @@
+---
+license: apache-2.0
+language:
+- ar
+base_model: oddadmix/50M-2048-Emhotob
+library_name: transformers
+tags: [arabic, zero-shot-classification, prompt-routing, router]
+---
+
+# Nawah-Router-v3 — موجّه عربي صفري
+
+**52M parameters.** Give it a text and **any categories in plain Arabic**; it scores all of them in
+**one forward pass**. Categories are chosen at inference — no fixed taxonomy.
+
+> **بالعربية:** نموذج عربي يوجّه أي نص إلى فئة من فئات تكتبها أنت بلغة طبيعية، في مسار واحد.
+
+## Results
+
+Trained on [`oddadmix/arabic-prompt-routing`](https://huggingface.co/datasets/oddadmix/arabic-prompt-routing)
+(233,720 rows, 12 routing axes).
+
+| eval | v3 | v2 (51K corpus) | random |
+|---|---:|---:|---:|
+| unseen category sets | **0.9305** | 0.9199 | 0.2137 |
+| unseen domains | **0.6976** | 0.6665 | 0.2521 |
+| **unseen axes** | **0.6130** | *(n/a)* | 0.2358 |
+| deliberately adjacent categories | 0.9008 | 0.9149 | 0.2109 |
+
+`unseen_axis` is the strongest claim here: `tools` and `retrieval` appear **nowhere** in training,
+and the model still routes along them at 0.61 against a 0.24 baseline.
+
+## Usage
+
+```python
+from transformers import AutoTokenizer
+from routing_model import RouterModel, route      # ships in this repo
+
+M = "oddadmix/Nawah-Router-v3"
+tok = AutoTokenizer.from_pretrained(M)
+model = RouterModel.from_pretrained(M)
+
+route(model, tok, "كم صار سعر صرف الدولار اليوم؟",
+      ["بحث في الويب", "حاسبة", "تقويم ومواعيد", "لا يحتاج أداة"])
+```
+
+## Choosing a checkpoint
+
+No single configuration wins everything, and the trade-off is real:
+
+| epochs | unseen lanes | unseen domains | unseen axes |
+|---:|---:|---:|---:|
+| 1 | **0.9327** | 0.6968 | 0.5950 |
+| **2** (this) | 0.9305 | **0.6976** | 0.6130 |
+| 3 | 0.9181 | 0.6496 | **0.6376** |
+
+Longer training helps the hardest transfer (unseen axes) and costs the everyday cases. 2 epochs is
+shipped as the balance. Higher learning rates are simply worse — 6e-4 and 1e-3 both degrade.
+
+## How the head works
+
+Text and categories share one sequence, **text first**. Each category's span is mean-pooled into
+its own vector and a **shared** scorer turns each into one logit; the softmax runs over the
+categories supplied. Because the scorer is shared it reads category *content*, not slot index —
+which is what makes the label set free text.
+
+A fixed-slot head (`num_labels = max_lanes`) scored **exactly 1/n at every lane count**: its
+weights were positional, and the corpus randomises category order, so there was nothing to learn.
+
+## Limitations
+
+Unseen domains (0.70) and unseen axes (0.61) trail unseen category sets (0.93) — it generalises
+best inside verticals and dimensions it has seen. Confidence is **not calibrated**: clear cases
+saturate near 1.0, so use the ranking, not the number. Arabic, 1–3 line messages.
+
+3 epochs at LR 3e-4 cosine, batch 32, bf16, `max_length` 320.

build_router_dataset_v2.py

@@ -0,0 +1,164 @@
+"""
+Merge the v1 and v2 router caches, shuffle lanes, audit for shortcuts, and write the splits.
+
+v1 (out_router, task ids 100000-116063) is topical routing across 18 business verticals.
+v2 (out_router_v2, ids 1000000+) adds compound lane sets, ten single-axis taxonomies and 45
+verticals. The id ranges are disjoint by construction, so the two caches concatenate safely and
+v1 never had to be regenerated.
+
+Four splits, in increasing order of how much they ask of the model:
+
+  unseen_lanes  - new lane sets, axes and domains seen in training
+  unseen_domain - three business verticals held out entirely (kept identical to v1 so the 0.6665
+                  number stays comparable)
+  unseen_axis   - `tools` and `retrieval` never appear in training at all. The strongest
+                  zero-shot claim available: routing on a dimension the model was never taught.
+  hard          - deliberately adjacent lanes
+
+Lanes are shuffled here and the label recomputed, so whatever positional bias the generator had
+cannot survive into training.
+"""
+import json
+import os
+import random
+from collections import Counter, defaultdict
+from pathlib import Path
+
+import router_axes as ra
+
+V1 = Path("out_router"); V2 = Path("out_router_v2")
+OUT = Path(os.environ.get("DATA_DIR", "data_v2")); OUT.mkdir(exist_ok=True)
+HELD_DOMAINS = ["فرز طبي أولي", "استشارات قانونية", "شركة تأمين"]
+HELD_AXES = ["tools", "retrieval"]
+EVAL_TASK_FRAC = float(os.environ.get("EVAL_TASK_FRAC", 0.02))
+SEED = 42
+
+
+def axes_map(cache):
+    m = {}
+    p = cache / "generations.jsonl"
+    if not p.exists():
+        return m
+    for line in open(p, encoding="utf-8"):
+        try:
+            r = json.loads(line)
+        except json.JSONDecodeError:
+            continue
+        m[r["task_id"]] = r["axes"]
+    return m
+
+
+def load(cache, amap, default_axis):
+    kept, amb = [], 0
+    p = cache / "verified.jsonl"
+    if not p.exists():
+        return kept, amb
+    for line in open(p, encoding="utf-8"):
+        try:
+            r = json.loads(line)
+        except json.JSONDecodeError:
+            continue
+        if not r["agree"]:
+            amb += 1
+            continue
+        a = amap.get(r["task_id"], {})
+        axis = a.get("axis", default_axis)
+        r["axis"] = axis
+        r["family"] = axis.split(":")[0] if ":" in axis else axis
+        r["mode"] = a.get("mode", "topical")
+        kept.append(r)
+    return kept, amb
+
+
+def main():
+    rng = random.Random(SEED)
+    a1, a2 = axes_map(V1), axes_map(V2)
+    r1, amb1 = load(V1, a1, "topical")
+    r2, amb2 = load(V2, a2, "topical")
+    print(f"v1 kept {len(r1):,} (ambiguous {amb1:,}) | v2 kept {len(r2):,} (ambiguous {amb2:,})")
+
+    rows = []
+    for r in r1 + r2:
+        order = list(range(len(r["routes"])))
+        rng.shuffle(order)
+        rows.append({"task_id": r["task_id"], "text": r["text"],
+                     "routes": [r["routes"][i] for i in order],
+                     "label": order.index(r["label"]),
+                     "axis": r["axis"], "family": r["family"], "mode": r["mode"],
+                     "domain": r.get("domain", ""), "style": r.get("style", "phrase"),
+                     "difficulty": r.get("difficulty", "easy"),
+                     "has_other": r.get("has_other", False)})
+
+    seen, ded = set(), []
+    for r in rows:
+        k = ra.dedup_key(r["text"])
+        if k in seen:
+            continue
+        seen.add(k); ded.append(r)
+    print(f"merged {len(rows):,} -> {len(ded):,} after dedup ({len(rows)-len(ded):,} duplicates)")
+    rows = ded
+
+    axis_rows = [r for r in rows if r["family"] in HELD_AXES]
+    rest = [r for r in rows if r["family"] not in HELD_AXES]
+    dom_rows = [r for r in rest if r["domain"] in HELD_DOMAINS]
+    rest = [r for r in rest if r["domain"] not in HELD_DOMAINS]
+    tasks = sorted({r["task_id"] for r in rest}); rng.shuffle(tasks)
+    ev_tasks = set(tasks[: max(1, int(len(tasks) * EVAL_TASK_FRAC))])
+    unseen_lanes = [r for r in rest if r["task_id"] in ev_tasks]
+    train = [r for r in rest if r["task_id"] not in ev_tasks]
+
+    # A lane set invented twice by different tasks would put a vocabulary the model trained on
+    # into a set labelled "unseen". Drop those from eval rather than weaken the claim.
+    tls = {ra.lane_set_key(r["routes"]) for r in train}
+    before = len(unseen_lanes), len(dom_rows), len(axis_rows)
+    unseen_lanes = [r for r in unseen_lanes if ra.lane_set_key(r["routes"]) not in tls]
+    dom_rows = [r for r in dom_rows if ra.lane_set_key(r["routes"]) not in tls]
+    axis_rows = [r for r in axis_rows if ra.lane_set_key(r["routes"]) not in tls]
+    print(f"lane-set overlap dropped: lanes {before[0]-len(unseen_lanes)}, "
+          f"domain {before[1]-len(dom_rows)}, axis {before[2]-len(axis_rows)}")
+    hard = [r for r in unseen_lanes if r["difficulty"] == "hard"]
+
+    for name, s in (("train", train), ("eval_unseen_lanes", unseen_lanes),
+                    ("eval_unseen_domain", dom_rows), ("eval_unseen_axis", axis_rows),
+                    ("eval_hard", hard)):
+        with open(OUT / f"{name}.jsonl", "w", encoding="utf-8") as fh:
+            for r in s:
+                fh.write(json.dumps(r, ensure_ascii=False) + "\n")
+        print(f"  {name:<20}{len(s):>8,}")
+
+    leaked = {r["family"] for r in train} & set(HELD_AXES)
+    print(f"\nheld-out axes present in train: {leaked or 'none'}  (must be none)")
+    print(f"train texts also in any eval: "
+          f"{len({r['text'] for r in train} & {r['text'] for r in unseen_lanes+dom_rows+axis_rows})}")
+
+    print("\n--- composition ---")
+    print("by mode:", dict(Counter(r["mode"] for r in rows).most_common()))
+    fam = Counter(r["family"] for r in rows)
+    print("by axis family:", dict(fam.most_common()))
+    print(f"distinct lane sets: {len({ra.lane_set_key(r['routes']) for r in rows}):,}")
+
+    print("\n--- shortcut audit ---")
+    groups = defaultdict(list)
+    for r in rows:
+        groups[len(r["routes"])].append(r)
+    for k in sorted(groups):
+        g = groups[k]
+        if len(g) < 200:
+            continue
+        c = Counter(r["label"] for r in g); exp = len(g) / k
+        chi = sum((c[i] - exp) ** 2 / exp for i in range(k))
+        crit = {1: 3.8, 2: 6.0, 3: 7.8, 4: 9.5, 5: 11.1, 6: 12.6, 7: 14.1, 8: 15.5}.get(k - 1, 15.5)
+        print(f"  {k} lanes n={len(g):<7} chi2={chi:7.1f} crit={crit:<5} "
+              f"{'FLAT' if chi < crit else 'SKEWED'}")
+    ranks = []
+    for r in rows:
+        L = [len(x) for x in r["routes"]]
+        ranks.append(sorted(L, reverse=True).index(L[r["label"]]) / max(len(L) - 1, 1))
+    print(f"  correct-lane length rank: {sum(ranks)/len(ranks):.3f} (0.5 = no bias)")
+    give = sum(1 for r in rows if r["routes"][r["label"]] in r["text"])
+    print(f"  verbatim lane in prompt: {give} ({give/len(rows):.2%})")
+    print(f"  TOTAL ROWS: {len(rows):,}")
+
+
+if __name__ == "__main__":
+    main()

config.json

@@ -0,0 +1,32 @@
+{
+  "architectures": [
+    "LlamaModel"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "bos_token_id": 0,
+  "dtype": "float32",
+  "eos_token_id": 2,
+  "head_dim": 64,
+  "hidden_act": "silu",
+  "hidden_size": 512,
+  "initializer_range": 0.02,
+  "intermediate_size": 1408,
+  "max_position_embeddings": 2048,
+  "mlp_bias": false,
+  "model_type": "llama",
+  "num_attention_heads": 8,
+  "num_hidden_layers": 12,
+  "num_key_value_heads": 4,
+  "pad_token_id": 1,
+  "pretraining_tp": 1,
+  "rms_norm_eps": 1e-06,
+  "rope_parameters": {
+    "rope_theta": 10000,
+    "rope_type": "default"
+  },
+  "tie_word_embeddings": true,
+  "transformers_version": "5.15.1",
+  "use_cache": false,
+  "vocab_size": 32000
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b5200d894303508b35530632e7f86e433596cb9dd9a0e6eb124aaf20a75b8f8c
+size 207156480

router_axes.py

@@ -0,0 +1,460 @@
+"""
+Arabic prompt-routing corpus v2 — many routing axes, and lane sets that mix them.
+
+v1 (router_common.py) generated one axis only: topical intent inside a business vertical, 18
+domains. It reached 0.9199 on unseen lane sets but stalled at 0.6665 on unseen domains, and the
+reference demo shows why that is the wrong shape rather than merely too small. LiquidAI's shipped
+presets mix axes inside a single lane set:
+
+    Simple function call · Simple tool use · Complex multi-step agentic task ·
+    Quick factual question · Casual conversation · Creative writing · Translation ·
+    Needs a bigger model
+
+That is complexity, task type, tool use and an escalation lane together. Every v1 lane set was
+single-axis, so the model never saw categories that differ along different dimensions at once.
+
+Three generation modes here:
+
+  topical  - v1's mode, widened from 18 to 45 verticals. The model invents the lanes.
+  axis     - lanes drawn from one of ten taxonomies (task type, complexity, safety, retrieval,
+             tools, subject, subdomain, language, urgency, tone), each with several paraphrased
+             wordings so the model learns the concept and not the string.
+  mixed    - lanes sampled across two or three axes, as the reference does. The largest share.
+
+`router_common.py` is deliberately untouched: it defines build_task for task ids 100000-116063 and
+editing it would silently redraw prompts for the 51,392 rows already generated. This module owns a
+disjoint range (START_TASK=1000000) and its output is merged with v1's at build time.
+
+Labels are still established by construction and confirmed by the independent back-classification
+in verify_router.py, which is axis-agnostic and needs no change.
+"""
+import random
+import re
+import unicodedata
+
+ITEMS_PER_TASK = 4
+
+C, Q, A, E = "### فئات", "### نص", "### الفئة", "### نهاية"
+
+# Varied on purpose. v1 used one fixed string, which teaches the token rather than the concept;
+# the reference ships a catch-all in most presets under several different names.
+OTHER_LANES = ["أخرى / لا ينطبق", "غير ذلك", "لا ينطبق أي مما سبق", "أخرى",
+               "خارج التصنيفات المذكورة"]
+
+# ---------------------------------------------------------------- axis taxonomies
+# Each lane carries several wordings; build_task picks one per task so the corpus never ties a
+# concept to a single surface form. Wordings deliberately avoid a shared give-away token
+# (e.g. not every complexity lane contains بسيط/معقد).
+
+AXES = {
+    "task_type": {
+        "hint": "مستخدم يطلب من مساعد ذكي مهمة",
+        "lanes": [
+            ["ترجمة", "نقل نص من لغة إلى أخرى", "طلب ترجمة"],
+            ["برمجة", "كتابة أو إصلاح كود", "مهمة برمجية"],
+            ["تلخيص", "اختصار نص طويل", "طلب ملخّص"],
+            ["كتابة إبداعية", "تأليف قصة أو نص أدبي", "صياغة محتوى إبداعي"],
+            ["رياضيات", "حساب أو مسألة رياضية", "عملية حسابية"],
+            ["استخراج بيانات", "انتزاع حقول من نص", "تحويل نص إلى بيانات منظمة"],
+            ["سؤال معلومات", "طلب معلومة مباشرة", "سؤال معرفي عام"],
+            ["محادثة", "دردشة غير مهمّية", "كلام عابر"],
+            ["إعادة صياغة", "تحسين صياغة نص قائم", "تحرير وتنقيح"],
+            ["تصنيف", "تحديد فئة نص", "فرز محتوى"],
+        ],
+    },
+    "complexity": {
+        "hint": "مستخدم يكتب طلبًا يتفاوت في صعوبته",
+        "lanes": [
+            ["طلب بسيط ومباشر", "يُنجز بخطوة واحدة", "مهمة سهلة لا تحتاج تفكيرًا"],
+            ["طلب متوسط", "يحتاج خطوتين أو ثلاثًا", "مهمة معتدلة"],
+            ["طلب معقّد متعدد الخطوات", "يحتاج تخطيطًا وتسلسلًا طويلًا", "مهمة مركّبة"],
+            ["يحتاج نموذجًا أكبر", "يفوق قدرة نموذج صغير", "يُحال إلى نموذج أقوى"],
+            ["يحتاج استدلالًا خطوة بخطوة", "يتطلب سلسلة تفكير", "حل يحتاج برهنة"],
+        ],
+    },
+    "safety": {
+        "hint": "رسالة واردة تحتاج فرزًا من ناحية السلامة",
+        "lanes": [
+            ["محتوى آمن", "لا مشكلة فيه", "مقبول للنشر"],
+            ["رسائل مزعجة", "إعلانات متكررة غير مرغوبة", "سبام تجاري"],
+            ["يحتوي بيانات شخصية", "فيه معلومات تعريف حساسة", "يكشف خصوصية أحدهم"],
+            ["لغة مسيئة", "ألفاظ جارحة", "إهانة أو تنمّر"],
+        ],
+    },
+    "retrieval": {
+        "hint": "سؤال قد يحتاج مصدرًا خارجيًا أو لا",
+        "lanes": [
+            ["يُجاب من المعرفة العامة", "لا يحتاج بحثًا", "إجابة مباشرة بلا مصادر"],
+            ["يحتاج بحثًا واحدًا", "يكفيه مصدر واحد", "استعلام واحد يكفي"],
+            ["يحتاج بحثًا متعدد الخطوات", "يجمع بين أكثر من مصدر", "يتطلب تتبع عدة معلومات"],
+            ["يحتاج معلومة لحظية", "يعتمد على بيانات متغيرة الآن", "يحتاج تحديثًا آنيًا"],
+        ],
+    },
+    "tools": {
+        "hint": "طلب قد يستدعي أداة معينة",
+        "lanes": [
+            ["بحث في الويب", "استعلام محرك بحث", "تصفح الإنترنت"],
+            ["حاسبة", "إجراء عملية حسابية", "أداة حساب"],
+            ["تقويم ومواعيد", "جدولة موعد", "إدارة الأجندة"],
+            ["بريد ورسائل", "إرسال رسالة", "التعامل مع البريد"],
+            ["قاعدة بيانات", "استعلام من سجلات", "جلب صف من جدول"],
+            ["لا يحتاج أداة", "يُجاب بلا استدعاء خارجي", "بدون أدوات"],
+            ["تشغيل كود", "تنفيذ سكربت", "مفسّر برمجي"],
+        ],
+    },
+    "subject": {
+        "hint": "سؤال أكاديمي في تخصص ما",
+        "lanes": [
+            ["أحياء", "علوم الحياة", "بيولوجيا"],
+            ["قانون", "مسائل قانونية", "تشريع وأنظمة"],
+            ["اقتصاد", "مسائل اقتصادية", "اقتصاد وأسواق"],
+            ["هندسة", "مسائل هندسية", "تصميم وتنفيذ هندسي"],
+            ["طب وصحة", "مسائل طبية", "صحة وعلاج"],
+            ["تاريخ", "أحداث ووقائع تاريخية", "ماضٍ وسير"],
+            ["فلسفة", "مسائل فلسفية", "فكر ومنطق"],
+            ["فيزياء", "مسائل فيزيائية", "قوانين الطبيعة"],
+            ["حاسوب", "علوم الحاسب", "تقنية معلومات"],
+            ["علم نفس", "مسائل نفسية", "سلوك وإدراك"],
+        ],
+    },
+    "language": {
+        "hint": "رسالة مكتوبة بلهجة أو لغة ما",
+        "lanes": [
+            ["عربية فصحى", "لغة عربية رسمية", "فصحى معاصرة"],
+            ["لهجة مصرية", "عامية مصر", "مصري"],
+            ["لهجة خليجية", "عامية الخليج", "خليجي"],
+            ["لهجة شامية", "عامية بلاد الشام", "شامي"],
+            ["لهجة مغاربية", "عامية المغرب العربي", "مغاربي"],
+            ["إنجليزية", "نص بالإنجليزية", "لغة إنجليزية"],
+            ["عربي مختلط بالإنجليزية", "كتابة مزدوجة اللغة", "عربيزي أو تبديل لغوي"],
+        ],
+    },
+    "urgency": {
+        "hint": "رسالة تتفاوت في إلحاحها",
+        "lanes": [
+            ["عاجل", "يحتاج تدخلًا فوريًا", "أولوية قصوى"],
+            ["أولوية عادية", "يمكن معالجته ضمن الدور", "غير مستعجل"],
+            ["أولوية منخفضة", "يمكن تأجيله", "لا يستدعي استعجالًا"],
+        ],
+    },
+    "tone": {
+        "hint": "رسالة من عميل بنبرة معينة",
+        "lanes": [
+            ["شكوى غاضبة", "امتعاض شديد", "نبرة تصعيد"],
+            ["استفسار محايد", "سؤال هادئ", "نبرة عادية"],
+            ["شكر وثناء", "رسالة امتنان", "إطراء على الخدمة"],
+            ["طلب اعتذار أو تعويض", "مطالبة بحق", "طلب جبر ضرر"],
+        ],
+    },
+}
+
+# Mixed lane sets use COMPOUND lanes, not a union of two axes' lanes.
+#
+# The first attempt sampled lanes from two axes at once and produced ill-posed rows: an insulting
+# message written in Gulf dialect belongs to both 'لهجة خليجية' and 'إهانة أو تنمر', and the task
+# is single-label. Constraining to "compatible" axes did not fix it either - an angry complaint is
+# also urgent. Any two independent axes overlap by construction.
+#
+# The reference avoids this because its lanes are already compound: "Simple function call" fuses
+# complexity with tool use, "Complex multi-step agentic task" fuses complexity with task type.
+# Each lane names a point in the joint space, so the set partitions rather than overlaps. These
+# pools mirror that, one per facet-family; a lane set is drawn from a single pool.
+COMPOUND_POOLS = {
+    "assistant": (
+        "مستخدم يطلب من مساعد ذكي مهمة",
+        ["استدعاء دالة بسيط", "استخدام أداة بسيط", "مهمة وكيلة معقدة متعددة الخطوات",
+         "استدلال بسيط", "استدلال معقد يحتاج خطوات", "مهمة برمجية بسيطة",
+         "مهمة برمجية معقدة", "إصلاح خطأ برمجي صغير", "إعادة هيكلة كود واسعة",
+         "مراجعة كود", "كتابة اختبارات", "شرح كود قائم",
+         "مسألة حسابية بسيطة", "مسألة رياضية معقدة", "ترجمة", "تلخيص",
+         "كتابة إبداعية", "إعادة صياغة نص", "بحث عن معلومة مباشرة",
+         "سؤال سريع تُعرف إجابته", "محادثة عابرة", "يحتاج نموذجًا أكبر",
+         "يحتاج بحثًا في الويب", "يحتاج تنفيذ كود", "استعلام قاعدة بيانات",
+         "تحليل بيانات", "إنشاء رسم بياني", "استخراج حقول من نص",
+         "جدولة موعد في التقويم", "صياغة رسالة بريد"],
+    ),
+    "triage": (
+        "رسالة واردة إلى فريق خدمة العملاء",
+        ["شكوى غاضبة تحتاج تصعيدًا فوريًا", "استفسار عادي غير مستعجل",
+         "بلاغ عطل حرج يوقف الخدمة", "بلاغ عطل بسيط يمكن تأجيله",
+         "طلب استرجاع مبلغ", "طلب إلغاء اشتراك", "رسالة شكر وثناء",
+         "رسالة مزعجة إعلانية", "إساءة لفظية تستوجب حظرًا",
+         "طلب حذف بيانات شخصية", "استفسار عن الفاتورة", "طلب ميزة جديدة",
+         "مشكلة في الدخول إلى الحساب", "سؤال عام عن الخدمة"],
+    ),
+    "content": (
+        "نص يُراد تصنيف موضوعه ومستواه",
+        ["خبر سياسي", "خبر رياضي", "خبر اقتصادي", "تقرير علمي مبسّط",
+         "مقال رأي", "محتوى تعليمي للمبتدئين", "محتوى تقني متقدم",
+         "وصفة أو إرشادات عملية", "مراجعة منتج", "إعلان تجاري",
+         "نص قانوني رسمي", "نص طبي متخصص", "محتوى ترفيهي خفيف"],
+    ),
+}
+
+# "Which flavour of X" — the reference's `lang` preset generalised. The model invents the members.
+SUBDOMAIN_FAMILIES = [
+    ("لغة البرمجة المستخدمة في السؤال", "مبرمج يسأل عن لغة معينة"),
+    ("نوع الوثيقة الرسمية المطلوبة", "مواطن يطلب مستندًا"),
+    ("فرع التأمين المعني", "عميل يسأل عن وثيقته"),
+    ("نوع الحساب البنكي المقصود", "عميل بنك"),
+    ("تخصص الطبيب المطلوب", "مريض يبحث عن عيادة"),
+    ("نوع العقار المطلوب", "باحث عن سكن"),
+    ("نوع الشحنة المرسلة", "عميل شركة شحن"),
+    ("مرحلة دراسية", "ولي أمر يسأل عن التسجيل"),
+]
+
+# v1's 18, plus 27 more.
+DOMAINS = [
+    ("دعم فني لشركة برمجيات", "مستخدم يراسل الدعم الفني"),
+    ("خدمة عملاء بنك", "عميل يتواصل مع بنكه"),
+    ("متجر إلكتروني", "مشترٍ يسأل عن طلب أو منتج"),
+    ("حجز سفر وطيران", "مسافر يرتب رحلة"),
+    ("خدمات حكومية إلكترونية", "مواطن يطلب معاملة"),
+    ("فرز طبي أولي", "مريض يصف حالته"),
+    ("شركة اتصالات", "مشترك يسأل عن باقته"),
+    ("منصة تعليمية", "طالب أو ولي أمر يستفسر"),
+    ("موارد بشرية داخل شركة", "موظف يراسل الموارد البشرية"),
+    ("وساطة عقارية", "باحث عن سكن أو مالك"),
+    ("تطبيق توصيل طعام", "زبون يتابع طلبه"),
+    ("استشارات قانونية", "شخص يسأل عن حقوقه"),
+    ("تصنيف مواضيع الأخبار", "قارئ يبحث عن خبر"),
+    ("مراجعة المحتوى والإبلاغ", "مستخدم يبلّغ عن محتوى"),
+    ("توجيه مساعد ذكي", "مستخدم يطلب من مساعد ذكي مهمة"),
+    ("إدارة مالية شخصية", "شخص يدير مصروفاته"),
+    ("شركة تأمين", "عميل يسأل عن وثيقته"),
+    ("صيانة منزلية", "ساكن يطلب صيانة"),
+    ("مطعم وحجز طاولات", "زبون يحجز أو يشتكي"),
+    ("صيدلية وطلب دواء", "مريض يسأل عن دواء"),
+    ("نادٍ رياضي واشتراكات", "مشترك في صالة رياضية"),
+    ("شحن ولوجستيات", "عميل يتابع شحنة"),
+    ("معرض وصيانة سيارات", "مالك سيارة"),
+    ("بنك استثماري ومحافظ", "مستثمر يسأل عن محفظته"),
+    ("سياحة دينية وحج وعمرة", "معتمر ينظّم رحلته"),
+    ("جامعة وقبول وتسجيل", "طالب جامعي"),
+    ("مكتب محاماة", "موكّل يتابع قضيته"),
+    ("زراعة وإرشاد زراعي", "مزارع يسأل عن محصوله"),
+    ("خدمة كهرباء ومياه", "مشترك يبلّغ عن انقطاع"),
+    ("تأجير سيارات", "مستأجر يحجز مركبة"),
+    ("فندق وضيافة", "نزيل يحجز أو يشتكي"),
+    ("منصة توظيف", "باحث عن عمل"),
+    ("متجر إلكترونيات وضمان", "مشترٍ يطالب بالضمان"),
+    ("خدمات بريدية", "مرسِل يتابع طردًا"),
+    ("مركز تدريب ودورات", "متدرب يسجّل في دورة"),
+    ("عيادة أسنان ومواعيد", "مريض يحجز موعدًا"),
+    ("شركة تنظيف منازل", "زبون يطلب خدمة"),
+    ("منصة بث ومحتوى", "مشترك يسأل عن اشتراكه"),
+    ("متجر ملابس ومقاسات", "زبون يستبدل قطعة"),
+    ("خدمة استضافة ونطاقات", "صاحب موقع"),
+    ("جمعية خيرية وتبرعات", "متبرع يسأل"),
+    ("نقل عام وتذاكر", "راكب يستفسر"),
+    ("مكتبة وبيع كتب", "قارئ يطلب كتابًا"),
+    ("خدمة محاسبة وضرائب", "صاحب منشأة"),
+    ("متجر حيوانات أليفة", "مربّي حيوان"),
+]
+
+STYLES = {
+    "noun":   ("كلمة واحدة أو كلمتين", "مثل: برمجة، فواتير، شكاوى"),
+    "phrase": ("عبارة قصيرة من ٢ إلى ٤ كلمات", "مثل: كتابة إبداعية، إلغاء الاشتراك"),
+    "desc":   ("جملة وصفية كاملة", "مثل: أسئلة حول الفواتير والمدفوعات وطلبات الاسترجاع"),
+}
+
+DIFFICULTY = {
+    "easy": "اجعل الفئات متباعدة وواضحة الفروق بينها.",
+    "hard": ("اجعل بعض الفئات متقاربة جدًا في المعنى بحيث يصعب التمييز بينها دون قراءة النص "
+             "بعناية."),
+}
+
+# Mode 1: the model invents the lanes (topical / subdomain).
+PROMPT_INVENT = """أنت تكتب بيانات تدريبية لنموذج يوجّه النصوص إلى فئات يحددها المستخدم.
+
+المجال: {domain} ({hint})
+
+الخطوة الأولى — اقترح {n} فئة للتوجيه في هذا المجال:
+- اكتب كل فئة بأسلوب: {style_desc} ({style_ex}).
+- {difficulty}
+- الفئات يجب أن تغطي المجال دون تداخل كامل بينها.{other_note}
+
+الخطوة الثانية — اكتب {k} نصوص مستقلة، كل نص من شخص حقيقي ({hint}).
+النص رقم i يجب أن ينتمي إلى الفئة المطلوبة له في القائمة أدناه:
+{targets}
+
+{fmt}"""
+
+# Mode 2/3: the lanes are given (single axis or mixed axes); the model only writes the texts.
+PROMPT_GIVEN = """أنت تكتب بيانات تدريبية لنموذج يوجّه النصوص إلى فئات يحددها المستخدم.
+
+الفئات المعطاة (استخدمها كما هي حرفيًا وبنفس الترتيب):
+{lanes_block}
+
+السياق: {hint}
+
+اكتب {k} نصوص عربية مستقلة وواقعية. النص رقم i يجب أن ينتمي إلى الفئة المطلوبة له:
+{targets}
+
+{fmt}"""
+
+FMT = """اكتب بهذا الشكل بالضبط:
+
+{C}
+(فئة في كل سطر، مسبوقة بـ "- "، بالترتيب وبنفس الصياغة في كل مرة)
+{Q}
+(النص العربي الأول، من سطر إلى ثلاثة أسطر)
+{A}
+(اسم الفئة التي ينتمي إليها هذا النص، منسوخًا حرفيًا من قائمة الفئات)
+{Q}
+(النص الثاني)
+{A}
+(فئته)
+... وهكذا حتى {k} نصوص
+{E}
+
+قواعد مهمة:
+- لا تذكر اسم الفئة داخل النص نفسه. النص يجب أن يدل على فئته بالمعنى لا بالكلمة.
+- اكتب نصوصًا طبيعية كما يكتبها الناس فعلًا، بأطوال مختلفة.
+- نوّع الأسماء والتفاصيل والأسلوب بين النصوص.
+- لا تكتب أي شيء خارج الوسوم."""
+
+
+def _targets_block(rng, n_lanes, has_other):
+    real = list(range(n_lanes))
+    targets = [rng.choice(real) for _ in range(ITEMS_PER_TASK)]
+    if has_other:
+        targets[rng.randrange(ITEMS_PER_TASK)] = n_lanes
+    return targets, "\n".join(f"- النص {i+1}: الفئة رقم {t+1}"
+                              for i, t in enumerate(targets))
+
+
+def build_task(task_id: int, seed: int = 1234):
+    """task_id -> (axes, prompt). Deterministic, so a resumed run redraws identical prompts."""
+    rng = random.Random(seed * 7_000_003 + task_id)
+    mode = rng.choices(["mixed", "axis", "topical"], weights=[45, 35, 20])[0]
+    style = rng.choice(["noun", "phrase", "phrase", "desc"])
+    difficulty = "hard" if rng.random() < 0.4 else "easy"
+    has_other = rng.random() < 0.30            # v1 used 0.15; the reference uses a catch-all often
+    other_lane = rng.choice(OTHER_LANES)
+    fmt = FMT.format(C=C, Q=Q, A=A, E=E, k=ITEMS_PER_TASK)
+
+    if mode == "topical":
+        use_sub = rng.random() < 0.2
+        domain, hint = rng.choice(SUBDOMAIN_FAMILIES if use_sub else DOMAINS)
+        n = rng.choice([2, 3, 3, 4, 4, 4, 5, 5, 6, 7, 8])
+        sd, sx = STYLES[style]
+        other_note = (f"\n- أضف في نهاية القائمة فئة أخيرة اسمها بالضبط: {other_lane}"
+                      "\n- واجعل أحد النصوص لا ينتمي لأي فئة إلا هذه الفئة الأخيرة.") \
+            if has_other else ""
+        targets, tgt = _targets_block(rng, n, has_other)
+        axes = {"mode": mode, "axis": "subdomain" if use_sub else "topical", "domain": domain,
+                "style": style, "difficulty": difficulty, "has_other": has_other,
+                # total lane count, Other included - the axis/mixed branch counts the same way
+                "n_routes": n + (1 if has_other else 0), "targets": targets}
+        return axes, PROMPT_INVENT.format(domain=domain, hint=hint, n=n, k=ITEMS_PER_TASK,
+                                          style_desc=sd, style_ex=sx,
+                                          difficulty=DIFFICULTY[difficulty],
+                                          other_note=other_note, targets=tgt, fmt=fmt)
+
+    if mode == "axis":
+        name = rng.choice(list(AXES))
+        spec = AXES[name]
+        k = min(len(spec["lanes"]), rng.choice([3, 3, 4, 4, 5, 5, 6, 7]))
+        picked = rng.sample(spec["lanes"], k)
+        lanes = [rng.choice(w) for w in picked]        # one wording per lane, per task
+        hint = spec["hint"]
+        axis_label, domain = name, name
+    else:                                              # mixed - compound lanes from one pool
+        pool_name = rng.choice(list(COMPOUND_POOLS))
+        hint, pool = COMPOUND_POOLS[pool_name]
+        k = min(len(pool), rng.choice([3, 4, 4, 5, 5, 6, 6, 7, 8]))
+        lanes = rng.sample(pool, k)
+        axis_label, domain = "compound:" + pool_name, pool_name
+
+    if has_other:
+        lanes = lanes + [other_lane]
+    rng.shuffle(lanes)
+    n = len(lanes) - (1 if has_other else 0)
+    other_idx = lanes.index(other_lane) if has_other else None
+    real_idx = [i for i in range(len(lanes)) if i != other_idx]
+    targets = [rng.choice(real_idx) for _ in range(ITEMS_PER_TASK)]
+    if has_other:
+        targets[rng.randrange(ITEMS_PER_TASK)] = other_idx
+    tgt = "\n".join(f"- النص {i+1}: الفئة رقم {t+1}" for i, t in enumerate(targets))
+    lanes_block = "\n".join(f"- {c}" for c in lanes)
+    axes = {"mode": mode, "axis": axis_label, "domain": domain, "style": style,
+            "difficulty": difficulty, "has_other": has_other,
+            "n_routes": len(lanes), "targets": targets}
+    return axes, PROMPT_GIVEN.format(lanes_block=lanes_block, hint=hint, k=ITEMS_PER_TASK,
+                                     targets=tgt, fmt=fmt)
+
+
+# ---------------------------------------------------------------- parsing / validation
+# Identical contract to router_common; the tag format and structural checks carry over unchanged.
+
+BLOCK_RE = re.compile(re.escape(C) + r"(?P<cats>.*?)(?=" + re.escape(Q) + r")", re.S)
+PAIR_RE = re.compile(re.escape(Q) + r"(?P<text>.*?)" + re.escape(A) + r"(?P<label>.*?)"
+                     r"(?=" + re.escape(Q) + r"|" + re.escape(E) + r"|\Z)", re.S)
+
+
+def norm(s: str) -> str:
+    s = unicodedata.normalize("NFKC", s or "")
+    s = re.sub(r"[ً-ْٰـ]", "", s)
+    return re.sub(r"\s+", " ", s).strip()
+
+
+def parse_items(raw: str):
+    m = BLOCK_RE.search(raw or "")
+    if not m:
+        return []
+    routes = []
+    for line in m.group("cats").split("\n"):
+        line = line.strip()
+        c = re.match(r"^[-*•]\s*(.+)$", line) or re.match(r"^\d+[.)]\s*(.+)$", line)
+        if c:
+            r = norm(c.group(1)).strip(" .:،")
+            if 2 <= len(r) <= 120:
+                routes.append(r)
+    if not routes:
+        return []
+    out = []
+    for p in PAIR_RE.finditer(raw):
+        out.append({"routes": routes, "text": norm(p.group("text")),
+                    "label_text": norm(p.group("label")).strip(" .:،").lstrip("- ")})
+    return out
+
+
+def _match_route(label, routes):
+    if label in routes:
+        return routes.index(label)
+    hits = [i for i, r in enumerate(routes) if r and (r in label or label in r)]
+    return hits[0] if len(hits) == 1 else None
+
+
+def validate(item, min_text=15, max_text=600):
+    """Structural only. Semantic correctness is decided by verify_router.py, not here."""
+    routes, text, label = item["routes"], item["text"], item["label_text"]
+    if not (2 <= len(routes) <= 9):
+        return False, "route_count"
+    if len(set(routes)) != len(routes):
+        return False, "duplicate_routes"
+    if not (min_text <= len(text) <= max_text):
+        return False, "text_length"
+    if any(t in text for t in (C, Q, A, E)):
+        return False, "tag_leak"
+    idx = _match_route(label, routes)
+    if idx is None:
+        return False, "label_not_in_routes"
+    if routes[idx] and routes[idx] in text:
+        return False, "lexical_giveaway"
+    if len(text.split()) < 4:
+        return False, "text_too_few_words"
+    return True, "ok"
+
+
+def label_index(item):
+    return _match_route(item["label_text"], item["routes"])
+
+
+def dedup_key(text: str) -> str:
+    return re.sub(r"\d+", "#", re.sub(r"\W+", "", norm(text)))
+
+
+def lane_set_key(routes) -> str:
+    return "|".join(sorted(norm(r) for r in routes))

router_model.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:90ff8180f65ed7ae63a3445b5418bf3f35ea051816e6f7ff367450da44a7c0fc
+size 208245555

routing_model.py

@@ -0,0 +1,85 @@
+"""
+Inference for Nawah-Router-v2 — an Arabic zero-shot router with a routing head.
+
+The text and every category share one sequence. Each category's character span is mapped to token
+indices through the tokenizer's offset mapping and mean-pooled into its own vector; a shared
+scorer turns each into one logit, and the softmax runs over the categories actually supplied.
+Because the scorer is shared across positions it reads category *content*, not slot index — which
+is what makes the label set free text chosen at inference.
+
+Layout is text-first, categories-second on purpose. The backbone is causal, so this ordering is
+what lets every category token attend to the whole text; reversed, the categories would be
+encoded blind to it.
+"""
+import torch
+import torch.nn as nn
+from transformers import AutoModel, AutoTokenizer
+
+MAX_ROUTES = 9
+MAX_LENGTH = 320
+
+
+class RouterModel(nn.Module):
+    def __init__(self, base_model):
+        super().__init__()
+        self.backbone = AutoModel.from_pretrained(base_model, dtype=torch.float32)
+        h = self.backbone.config.hidden_size
+        self.score = nn.Sequential(nn.Linear(h, h), nn.GELU(), nn.Linear(h, 1))
+        self.config = self.backbone.config
+
+    def forward(self, input_ids, attention_mask, cat_pool, n_routes):
+        hs = self.backbone(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state
+        logits = self.score(torch.bmm(cat_pool.to(hs.dtype), hs)).squeeze(-1)
+        ar = torch.arange(logits.size(1), device=logits.device)[None, :]
+        return logits.masked_fill(ar >= n_routes[:, None], torch.finfo(logits.dtype).min)
+
+    @classmethod
+    def from_pretrained(cls, path, token=None):
+        import os
+        from huggingface_hub import hf_hub_download
+        m = cls(path)
+        w = (os.path.join(path, "router_model.pt") if os.path.isdir(path)
+             else hf_hub_download(path, "router_model.pt", token=token))
+        m.load_state_dict(torch.load(w, map_location="cpu", weights_only=True))
+        return m.eval()
+
+
+def build_text(text, routes):
+    head = f"النص:\n{text}\n\nالفئات:\n"
+    s, spans = head, []
+    for c in routes:
+        s += "- "
+        spans.append((len(s), len(s) + len(c)))
+        s += c + "\n"
+    return s, spans
+
+
+@torch.no_grad()
+def route(model, tok, text, routes):
+    """-> [{'route': str, 'score': float}] sorted high to low."""
+    routes = [r for r in routes if r and r.strip()][:MAX_ROUTES]
+    if not text.strip() or not routes:
+        return []
+    full, spans = build_text(text, routes)
+    enc = tok(full, return_offsets_mapping=True, add_special_tokens=False,
+              truncation=True, max_length=MAX_LENGTH)
+    ids, offs = enc["input_ids"], enc["offset_mapping"]
+    pool = torch.zeros(1, MAX_ROUTES, len(ids))
+    for ci, (s, e) in enumerate(spans):
+        idx = [t for t, (a, b) in enumerate(offs) if a < e and b > s and a != b]
+        if idx:
+            pool[0, ci, idx] = 1.0 / len(idx)
+    logits = model(torch.tensor([ids]), torch.ones(1, len(ids), dtype=torch.long),
+                   pool, torch.tensor([len(routes)]))
+    probs = logits.softmax(-1)[0][: len(routes)].tolist()
+    out = [{"route": r, "score": p} for r, p in zip(routes, probs)]
+    return sorted(out, key=lambda x: -x["score"])
+
+
+if __name__ == "__main__":
+    M = "oddadmix/Nawah-Router-v2"
+    tok = AutoTokenizer.from_pretrained(M)
+    model = RouterModel.from_pretrained(M)
+    for r in route(model, tok, "الطلب تأخر ساعة والسائق ما رد على الاتصال",
+                   ["استفسار عن التوصيل", "شكوى تأخير", "مشكلة في الدفع"]):
+        print(f"{r['score']:.3f}  {r['route']}")

tokenizer_config.json

@@ -0,0 +1,11 @@
+{
+  "backend": "tokenizers",
+  "bos_token": "<s>",
+  "eos_token": "</s>",
+  "is_local": false,
+  "local_files_only": false,
+  "model_max_length": 1000000000000000019884624838656,
+  "pad_token": "<pad>",
+  "tokenizer_class": "TokenizersBackend",
+  "unk_token": "<unk>"
+}

train_metrics.json

@@ -0,0 +1,1157 @@
+{
+  "results": {
+    "unseen_lanes": {
+      "accuracy": 0.9304995345950977,
+      "random_baseline": 0.2137126998901743,
+      "n": 3223,
+      "by_n_routes": {
+        "2": {
+          "acc": 0.954954954954955,
+          "n": 111
+        },
+        "4": {
+          "acc": 0.9449877750611247,
+          "n": 818
+        },
+        "5": {
+          "acc": 0.9423076923076923,
+          "n": 624
+        },
+        "8": {
+          "acc": 0.9341317365269461,
+          "n": 334
+        },
+        "3": {
+          "acc": 0.9104938271604939,
+          "n": 324
+        },
+        "7": {
+          "acc": 0.907928388746803,
+          "n": 391
+        },
+        "6": {
+          "acc": 0.9224952741020794,
+          "n": 529
+        },
+        "9": {
+          "acc": 0.8913043478260869,
+          "n": 92
+        }
+      },
+      "by_difficulty": {
+        "easy": {
+          "acc": 0.9478603049680275,
+          "n": 2033
+        },
+        "hard": {
+          "acc": 0.9008403361344538,
+          "n": 1190
+        }
+      },
+      "by_mode": {
+        "topical": {
+          "acc": 0.884,
+          "n": 1500
+        },
+        "mixed": {
+          "acc": 0.9772549019607844,
+          "n": 1275
+        },
+        "axis": {
+          "acc": 0.953125,
+          "n": 448
+        }
+      }
+    },
+    "unseen_domain": {
+      "accuracy": 0.6976068889556424,
+      "random_baseline": 0.25206113185236045,
+      "n": 9987,
+      "by_n_routes": {
+        "5": {
+          "acc": 0.6869671132764921,
+          "n": 1642
+        },
+        "4": {
+          "acc": 0.727150084317032,
+          "n": 2965
+        },
+        "3": {
+          "acc": 0.7427476737821566,
+          "n": 1827
+        },
+        "2": {
+          "acc": 0.7923497267759563,
+          "n": 915
+        },
+        "6": {
+          "acc": 0.6202393906420022,
+          "n": 919
+        },
+        "7": {
+          "acc": 0.5959595959595959,
+          "n": 792
+        },
+        "8": {
+          "acc": 0.5976608187134503,
+          "n": 855
+        },
+        "9": {
+          "acc": 0.6666666666666666,
+          "n": 72
+        }
+      },
+      "by_difficulty": {
+        "easy": {
+          "acc": 0.6926144520657203,
+          "n": 6269
+        },
+        "hard": {
+          "acc": 0.7060247444862829,
+          "n": 3718
+        }
+      },
+      "by_mode": {
+        "topical": {
+          "acc": 0.6976068889556424,
+          "n": 9987
+        }
+      }
+    },
+    "unseen_axis": {
+      "accuracy": 0.6129888942539836,
+      "random_baseline": 0.23578575644002975,
+      "n": 12426,
+      "by_n_routes": {
+        "4": {
+          "acc": 0.541494632535095,
+          "n": 4844
+        },
+        "5": {
+          "acc": 0.6220699603317706,
+          "n": 2773
+        },
+        "3": {
+          "acc": 0.6389937106918239,
+          "n": 2385
+        },
+        "6": {
+          "acc": 0.7419898819561551,
+          "n": 1186
+        },
+        "7": {
+          "acc": 0.7033898305084746,
+          "n": 944
+        },
+        "8": {
+          "acc": 0.6836734693877551,
+          "n": 294
+        }
+      },
+      "by_difficulty": {
+        "hard": {
+          "acc": 0.6019339242546333,
+          "n": 4964
+        },
+        "easy": {
+          "acc": 0.6203430715625837,
+          "n": 7462
+        }
+      },
+      "by_mode": {
+        "axis": {
+          "acc": 0.6129888942539836,
+          "n": 12426
+        }
+      }
+    },
+    "hard": {
+      "accuracy": 0.9008403361344538,
+      "random_baseline": 0.21085734293717487,
+      "n": 1190,
+      "by_n_routes": {
+        "4": {
+          "acc": 0.9182156133828996,
+          "n": 269
+        },
+        "5": {
+          "acc": 0.9043824701195219,
+          "n": 251
+        },
+        "8": {
+          "acc": 0.9291338582677166,
+          "n": 127
+        },
+        "3": {
+          "acc": 0.8571428571428571,
+          "n": 126
+        },
+        "7": {
+          "acc": 0.8662420382165605,
+          "n": 157
+        },
+        "6": {
+          "acc": 0.9045226130653267,
+          "n": 199
+        },
+        "2": {
+          "acc": 1.0,
+          "n": 34
+        },
+        "9": {
+          "acc": 0.8148148148148148,
+          "n": 27
+        }
+      },
+      "by_difficulty": {
+        "hard": {
+          "acc": 0.9008403361344538,
+          "n": 1190
+        }
+      },
+      "by_mode": {
+        "topical": {
+          "acc": 0.8300180831826401,
+          "n": 553
+        },
+        "mixed": {
+          "acc": 0.9680851063829787,
+          "n": 470
+        },
+        "axis": {
+          "acc": 0.9461077844311377,
+          "n": 167
+        }
+      }
+    }
+  },
+  "base_model": "oddadmix/50M-2048-Emhotob",
+  "formulation": "routing_head",
+  "log_history": [
+    {
+      "loss": 0.8975238800048828,
+      "grad_norm": 11.03390121459961,
+      "learning_rate": 7.4625e-05,
+      "epoch": 0.030902348578491966,
+      "step": 200
+    },
+    {
+      "loss": 0.47223701477050783,
+      "grad_norm": 11.903738975524902,
+      "learning_rate": 0.000149625,
+      "epoch": 0.06180469715698393,
+      "step": 400
+    },
+    {
+      "loss": 0.4622153091430664,
+      "grad_norm": 6.198873996734619,
+      "learning_rate": 0.000224625,
+      "epoch": 0.09270704573547589,
+      "step": 600
+    },
+    {
+      "loss": 0.49717689514160157,
+      "grad_norm": 4.740701198577881,
+      "learning_rate": 0.000299625,
+      "epoch": 0.12360939431396786,
+      "step": 800
+    },
+    {
+      "loss": 0.4837868118286133,
+      "grad_norm": 3.6331026554107666,
+      "learning_rate": 0.0002998012771863284,
+      "epoch": 0.15451174289245984,
+      "step": 1000
+    },
+    {
+      "eval_unseen_lanes_loss": 0.5073266625404358,
+      "eval_unseen_lanes_accuracy": 0.8163201985727583,
+      "eval_unseen_lanes_runtime": 0.5548,
+      "eval_unseen_lanes_samples_per_second": 5809.493,
+      "eval_unseen_lanes_steps_per_second": 91.928,
+      "epoch": 0.15451174289245984,
+      "step": 1000
+    },
+    {
+      "eval_unseen_domain_loss": 0.9687234163284302,
+      "eval_unseen_domain_accuracy": 0.6065885651346751,
+      "eval_unseen_domain_runtime": 1.5808,
+      "eval_unseen_domain_samples_per_second": 6317.746,
+      "eval_unseen_domain_steps_per_second": 99.318,
+      "epoch": 0.15451174289245984,
+      "step": 1000
+    },
+    {
+      "eval_unseen_axis_loss": 1.1325784921646118,
+      "eval_unseen_axis_accuracy": 0.5135200386286818,
+      "eval_unseen_axis_runtime": 1.2551,
+      "eval_unseen_axis_samples_per_second": 9900.252,
+      "eval_unseen_axis_steps_per_second": 155.364,
+      "epoch": 0.15451174289245984,
+      "step": 1000
+    },
+    {
+      "eval_hard_loss": 0.5369061231613159,
+      "eval_hard_accuracy": 0.807563025210084,
+      "eval_hard_runtime": 0.2574,
+      "eval_hard_samples_per_second": 4623.428,
+      "eval_hard_steps_per_second": 73.819,
+      "epoch": 0.15451174289245984,
+      "step": 1000
+    },
+    {
+      "loss": 0.4474765396118164,
+      "grad_norm": 9.281235694885254,
+      "learning_rate": 0.00029920164204065217,
+      "epoch": 0.18541409147095178,
+      "step": 1200
+    },
+    {
+      "loss": 0.4180959320068359,
+      "grad_norm": 6.990286827087402,
+      "learning_rate": 0.00029820269436406154,
+      "epoch": 0.21631644004944375,
+      "step": 1400
+    },
+    {
+      "loss": 0.4012123489379883,
+      "grad_norm": 2.00594425201416,
+      "learning_rate": 0.00029680710766818503,
+      "epoch": 0.24721878862793573,
+      "step": 1600
+    },
+    {
+      "loss": 0.39807247161865233,
+      "grad_norm": 4.734849452972412,
+      "learning_rate": 0.0002950186170007609,
+      "epoch": 0.27812113720642767,
+      "step": 1800
+    },
+    {
+      "loss": 0.37501571655273436,
+      "grad_norm": 2.752845048904419,
+      "learning_rate": 0.0002928420089494242,
+      "epoch": 0.30902348578491967,
+      "step": 2000
+    },
+    {
+      "eval_unseen_lanes_loss": 0.406538724899292,
+      "eval_unseen_lanes_accuracy": 0.8569655600372323,
+      "eval_unseen_lanes_runtime": 0.5556,
+      "eval_unseen_lanes_samples_per_second": 5801.001,
+      "eval_unseen_lanes_steps_per_second": 91.794,
+      "epoch": 0.30902348578491967,
+      "step": 2000
+    },
+    {
+      "eval_unseen_domain_loss": 1.1472351551055908,
+      "eval_unseen_domain_accuracy": 0.5958746370281366,
+      "eval_unseen_domain_runtime": 1.575,
+      "eval_unseen_domain_samples_per_second": 6340.999,
+      "eval_unseen_domain_steps_per_second": 99.683,
+      "epoch": 0.30902348578491967,
+      "step": 2000
+    },
+    {
+      "eval_unseen_axis_loss": 1.353710412979126,
+      "eval_unseen_axis_accuracy": 0.49806856591018833,
+      "eval_unseen_axis_runtime": 1.3077,
+      "eval_unseen_axis_samples_per_second": 9502.1,
+      "eval_unseen_axis_steps_per_second": 149.116,
+      "epoch": 0.30902348578491967,
+      "step": 2000
+    },
+    {
+      "eval_hard_loss": 0.4623798727989197,
+      "eval_hard_accuracy": 0.8378151260504202,
+      "eval_hard_runtime": 0.2506,
+      "eval_hard_samples_per_second": 4749.401,
+      "eval_hard_steps_per_second": 75.831,
+      "epoch": 0.30902348578491967,
+      "step": 2000
+    },
+    {
+      "loss": 0.3629868316650391,
+      "grad_norm": 3.957963228225708,
+      "learning_rate": 0.0002902831088312284,
+      "epoch": 0.3399258343634116,
+      "step": 2200
+    },
+    {
+      "loss": 0.3620159912109375,
+      "grad_norm": 2.070098638534546,
+      "learning_rate": 0.00028734876510218624,
+      "epoch": 0.37082818294190356,
+      "step": 2400
+    },
+    {
+      "loss": 0.33227527618408204,
+      "grad_norm": 3.895007371902466,
+      "learning_rate": 0.00028404683102855547,
+      "epoch": 0.40173053152039556,
+      "step": 2600
+    },
+    {
+      "loss": 0.3365111541748047,
+      "grad_norm": 3.198639154434204,
+      "learning_rate": 0.0002803861436689221,
+      "epoch": 0.4326328800988875,
+      "step": 2800
+    },
+    {
+      "loss": 0.31873382568359376,
+      "grad_norm": 2.5327401161193848,
+      "learning_rate": 0.0002763765002233335,
+      "epoch": 0.4635352286773795,
+      "step": 3000
+    },
+    {
+      "eval_unseen_lanes_loss": 0.34330183267593384,
+      "eval_unseen_lanes_accuracy": 0.8836487744337573,
+      "eval_unseen_lanes_runtime": 0.553,
+      "eval_unseen_lanes_samples_per_second": 5828.318,
+      "eval_unseen_lanes_steps_per_second": 92.226,
+      "epoch": 0.4635352286773795,
+      "step": 3000
+    },
+    {
+      "eval_unseen_domain_loss": 0.9895414113998413,
+      "eval_unseen_domain_accuracy": 0.6378291779313107,
+      "eval_unseen_domain_runtime": 1.5853,
+      "eval_unseen_domain_samples_per_second": 6299.774,
+      "eval_unseen_domain_steps_per_second": 99.035,
+      "epoch": 0.4635352286773795,
+      "step": 3000
+    },
+    {
+      "eval_unseen_axis_loss": 1.0701825618743896,
+      "eval_unseen_axis_accuracy": 0.5737164010944793,
+      "eval_unseen_axis_runtime": 1.2712,
+      "eval_unseen_axis_samples_per_second": 9774.846,
+      "eval_unseen_axis_steps_per_second": 153.396,
+      "epoch": 0.4635352286773795,
+      "step": 3000
+    },
+    {
+      "eval_hard_loss": 0.38013485074043274,
+      "eval_hard_accuracy": 0.873109243697479,
+      "eval_hard_runtime": 0.2542,
+      "eval_hard_samples_per_second": 4681.85,
+      "eval_hard_steps_per_second": 74.752,
+      "epoch": 0.4635352286773795,
+      "step": 3000
+    },
+    {
+      "loss": 0.29870388031005857,
+      "grad_norm": 2.18878173828125,
+      "learning_rate": 0.00027202863181277745,
+      "epoch": 0.49443757725587145,
+      "step": 3200
+    },
+    {
+      "loss": 0.30184398651123046,
+      "grad_norm": 3.0138494968414307,
+      "learning_rate": 0.00026735417475918285,
+      "epoch": 0.5253399258343634,
+      "step": 3400
+    },
+    {
+      "loss": 0.32383251190185547,
+      "grad_norm": 7.008852005004883,
+      "learning_rate": 0.00026236563944280554,
+      "epoch": 0.5562422744128553,
+      "step": 3600
+    },
+    {
+      "loss": 0.3035639762878418,
+      "grad_norm": 2.690004825592041,
+      "learning_rate": 0.0002570763768203474,
+      "epoch": 0.5871446229913473,
+      "step": 3800
+    },
+    {
+      "loss": 0.2934083366394043,
+      "grad_norm": 1.9328616857528687,
+      "learning_rate": 0.0002515005426934173,
+      "epoch": 0.6180469715698393,
+      "step": 4000
+    },
+    {
+      "eval_unseen_lanes_loss": 0.30602237582206726,
+      "eval_unseen_lanes_accuracy": 0.8920260626745269,
+      "eval_unseen_lanes_runtime": 0.5511,
+      "eval_unseen_lanes_samples_per_second": 5848.268,
+      "eval_unseen_lanes_steps_per_second": 92.542,
+      "epoch": 0.6180469715698393,
+      "step": 4000
+    },
+    {
+      "eval_unseen_domain_loss": 0.983905553817749,
+      "eval_unseen_domain_accuracy": 0.6368278762391109,
+      "eval_unseen_domain_runtime": 1.5835,
+      "eval_unseen_domain_samples_per_second": 6306.791,
+      "eval_unseen_domain_steps_per_second": 99.146,
+      "epoch": 0.6180469715698393,
+      "step": 4000
+    },
+    {
+      "eval_unseen_axis_loss": 1.2108259201049805,
+      "eval_unseen_axis_accuracy": 0.5799130854659584,
+      "eval_unseen_axis_runtime": 1.2562,
+      "eval_unseen_axis_samples_per_second": 9891.534,
+      "eval_unseen_axis_steps_per_second": 155.227,
+      "epoch": 0.6180469715698393,
+      "step": 4000
+    },
+    {
+      "eval_hard_loss": 0.3655548691749573,
+      "eval_hard_accuracy": 0.8756302521008403,
+      "eval_hard_runtime": 0.2592,
+      "eval_hard_samples_per_second": 4591.918,
+      "eval_hard_steps_per_second": 73.316,
+      "epoch": 0.6180469715698393,
+      "step": 4000
+    },
+    {
+      "loss": 0.2829387664794922,
+      "grad_norm": 2.094444513320923,
+      "learning_rate": 0.00024565305982296296,
+      "epoch": 0.6489493201483313,
+      "step": 4200
+    },
+    {
+      "loss": 0.2714531135559082,
+      "grad_norm": 2.0051519870758057,
+      "learning_rate": 0.00023954957799106823,
+      "epoch": 0.6798516687268232,
+      "step": 4400
+    },
+    {
+      "loss": 0.26623159408569336,
+      "grad_norm": 2.2412750720977783,
+      "learning_rate": 0.00023320643211700308,
+      "epoch": 0.7107540173053152,
+      "step": 4600
+    },
+    {
+      "loss": 0.2522362518310547,
+      "grad_norm": 1.7817517518997192,
+      "learning_rate": 0.000226640598539622,
+      "epoch": 0.7416563658838071,
+      "step": 4800
+    },
+    {
+      "loss": 0.2480511474609375,
+      "grad_norm": 5.033050537109375,
+      "learning_rate": 0.00021986964958311316,
+      "epoch": 0.7725587144622992,
+      "step": 5000
+    },
+    {
+      "eval_unseen_lanes_loss": 0.27606716752052307,
+      "eval_unseen_lanes_accuracy": 0.9053676698727893,
+      "eval_unseen_lanes_runtime": 0.552,
+      "eval_unseen_lanes_samples_per_second": 5838.329,
+      "eval_unseen_lanes_steps_per_second": 92.384,
+      "epoch": 0.7725587144622992,
+      "step": 5000
+    },
+    {
+      "eval_unseen_domain_loss": 0.96076899766922,
+      "eval_unseen_domain_accuracy": 0.6541503955141684,
+      "eval_unseen_domain_runtime": 1.5806,
+      "eval_unseen_domain_samples_per_second": 6318.291,
+      "eval_unseen_domain_steps_per_second": 99.326,
+      "epoch": 0.7725587144622992,
+      "step": 5000
+    },
+    {
+      "eval_unseen_axis_loss": 1.4210246801376343,
+      "eval_unseen_axis_accuracy": 0.5544825366167713,
+      "eval_unseen_axis_runtime": 1.2599,
+      "eval_unseen_axis_samples_per_second": 9862.704,
+      "eval_unseen_axis_steps_per_second": 154.774,
+      "epoch": 0.7725587144622992,
+      "step": 5000
+    },
+    {
+      "eval_hard_loss": 0.3613678812980652,
+      "eval_hard_accuracy": 0.8857142857142857,
+      "eval_hard_runtime": 0.2606,
+      "eval_hard_samples_per_second": 4565.623,
+      "eval_hard_steps_per_second": 72.897,
+      "epoch": 0.7725587144622992,
+      "step": 5000
+    },
+    {
+      "loss": 0.24242013931274414,
+      "grad_norm": 2.4511566162109375,
+      "learning_rate": 0.00021291170652769533,
+      "epoch": 0.8034610630407911,
+      "step": 5200
+    },
+    {
+      "loss": 0.2558748817443848,
+      "grad_norm": 1.6439733505249023,
+      "learning_rate": 0.00020578539111112834,
+      "epoch": 0.8343634116192831,
+      "step": 5400
+    },
+    {
+      "loss": 0.24663785934448243,
+      "grad_norm": 3.4409291744232178,
+      "learning_rate": 0.00019850977569083514,
+      "epoch": 0.865265760197775,
+      "step": 5600
+    },
+    {
+      "loss": 0.24483919143676758,
+      "grad_norm": 3.071554183959961,
+      "learning_rate": 0.00019110433220001754,
+      "epoch": 0.896168108776267,
+      "step": 5800
+    },
+    {
+      "loss": 0.23662748336791992,
+      "grad_norm": 2.491429567337036,
+      "learning_rate": 0.0001835888800343768,
+      "epoch": 0.927070457354759,
+      "step": 6000
+    },
+    {
+      "eval_unseen_lanes_loss": 0.2571614384651184,
+      "eval_unseen_lanes_accuracy": 0.9081600992863792,
+      "eval_unseen_lanes_runtime": 0.5519,
+      "eval_unseen_lanes_samples_per_second": 5839.58,
+      "eval_unseen_lanes_steps_per_second": 92.404,
+      "epoch": 0.927070457354759,
+      "step": 6000
+    },
+    {
+      "eval_unseen_domain_loss": 0.9326336979866028,
+      "eval_unseen_domain_accuracy": 0.6525483128066486,
+      "eval_unseen_domain_runtime": 1.5742,
+      "eval_unseen_domain_samples_per_second": 6344.232,
+      "eval_unseen_domain_steps_per_second": 99.734,
+      "epoch": 0.927070457354759,
+      "step": 6000
+    },
+    {
+      "eval_unseen_axis_loss": 1.3221510648727417,
+      "eval_unseen_axis_accuracy": 0.5856269113149847,
+      "eval_unseen_axis_runtime": 1.2492,
+      "eval_unseen_axis_samples_per_second": 9946.782,
+      "eval_unseen_axis_steps_per_second": 156.094,
+      "epoch": 0.927070457354759,
+      "step": 6000
+    },
+    {
+      "eval_hard_loss": 0.33343571424484253,
+      "eval_hard_accuracy": 0.888235294117647,
+      "eval_hard_runtime": 0.2587,
+      "eval_hard_samples_per_second": 4600.467,
+      "eval_hard_steps_per_second": 73.453,
+      "epoch": 0.927070457354759,
+      "step": 6000
+    },
+    {
+      "loss": 0.23579851150512696,
+      "grad_norm": 2.5624895095825195,
+      "learning_rate": 0.00017598353300890969,
+      "epoch": 0.957972805933251,
+      "step": 6200
+    },
+    {
+      "loss": 0.21076148986816406,
+      "grad_norm": 0.9841051697731018,
+      "learning_rate": 0.00016830864552674267,
+      "epoch": 0.9888751545117429,
+      "step": 6400
+    },
+    {
+      "loss": 0.17771633148193358,
+      "grad_norm": 0.9096925258636475,
+      "learning_rate": 0.00016058475810407297,
+      "epoch": 1.019777503090235,
+      "step": 6600
+    },
+    {
+      "loss": 0.14182785034179687,
+      "grad_norm": 3.3751485347747803,
+      "learning_rate": 0.00015283254239701005,
+      "epoch": 1.0506798516687268,
+      "step": 6800
+    },
+    {
+      "loss": 0.1347336483001709,
+      "grad_norm": 0.5445594787597656,
+      "learning_rate": 0.00014507274587744398,
+      "epoch": 1.0815822002472189,
+      "step": 7000
+    },
+    {
+      "eval_unseen_lanes_loss": 0.2655007839202881,
+      "eval_unseen_lanes_accuracy": 0.9162271175923054,
+      "eval_unseen_lanes_runtime": 0.5535,
+      "eval_unseen_lanes_samples_per_second": 5823.011,
+      "eval_unseen_lanes_steps_per_second": 92.142,
+      "epoch": 1.0815822002472189,
+      "step": 7000
+    },
+    {
+      "eval_unseen_domain_loss": 1.0792171955108643,
+      "eval_unseen_domain_accuracy": 0.6691699208971663,
+      "eval_unseen_domain_runtime": 1.5827,
+      "eval_unseen_domain_samples_per_second": 6310.28,
+      "eval_unseen_domain_steps_per_second": 99.2,
+      "epoch": 1.0815822002472189,
+      "step": 7000
+    },
+    {
+      "eval_unseen_axis_loss": 1.5451632738113403,
+      "eval_unseen_axis_accuracy": 0.5832930951231289,
+      "eval_unseen_axis_runtime": 1.2555,
+      "eval_unseen_axis_samples_per_second": 9897.184,
+      "eval_unseen_axis_steps_per_second": 155.316,
+      "epoch": 1.0815822002472189,
+      "step": 7000
+    },
+    {
+      "eval_hard_loss": 0.3464004695415497,
+      "eval_hard_accuracy": 0.8966386554621849,
+      "eval_hard_runtime": 0.2591,
+      "eval_hard_samples_per_second": 4593.088,
+      "eval_hard_steps_per_second": 73.335,
+      "epoch": 1.0815822002472189,
+      "step": 7000
+    },
+    {
+      "loss": 0.13339081764221192,
+      "grad_norm": 2.522960901260376,
+      "learning_rate": 0.00013732613630600555,
+      "epoch": 1.1124845488257107,
+      "step": 7200
+    },
+    {
+      "loss": 0.11780735015869141,
+      "grad_norm": 1.1648684740066528,
+      "learning_rate": 0.0001296134461507269,
+      "epoch": 1.1433868974042027,
+      "step": 7400
+    },
+    {
+      "loss": 0.13406469345092772,
+      "grad_norm": 1.038225531578064,
+      "learning_rate": 0.00012195531710015758,
+      "epoch": 1.1742892459826946,
+      "step": 7600
+    },
+    {
+      "loss": 0.1309179973602295,
+      "grad_norm": 1.6545023918151855,
+      "learning_rate": 0.00011437224481943516,
+      "epoch": 1.2051915945611866,
+      "step": 7800
+    },
+    {
+      "loss": 0.13036264419555665,
+      "grad_norm": 1.7858803272247314,
+      "learning_rate": 0.00010688452409716325,
+      "epoch": 1.2360939431396787,
+      "step": 8000
+    },
+    {
+      "eval_unseen_lanes_loss": 0.26563557982444763,
+      "eval_unseen_lanes_accuracy": 0.9162271175923054,
+      "eval_unseen_lanes_runtime": 0.5528,
+      "eval_unseen_lanes_samples_per_second": 5830.741,
+      "eval_unseen_lanes_steps_per_second": 92.264,
+      "epoch": 1.2360939431396787,
+      "step": 8000
+    },
+    {
+      "eval_unseen_domain_loss": 1.1347131729125977,
+      "eval_unseen_domain_accuracy": 0.6633623710824071,
+      "eval_unseen_domain_runtime": 1.5784,
+      "eval_unseen_domain_samples_per_second": 6327.271,
+      "eval_unseen_domain_steps_per_second": 99.467,
+      "epoch": 1.2360939431396787,
+      "step": 8000
+    },
+    {
+      "eval_unseen_axis_loss": 1.7664040327072144,
+      "eval_unseen_axis_accuracy": 0.5963302752293578,
+      "eval_unseen_axis_runtime": 1.2484,
+      "eval_unseen_axis_samples_per_second": 9953.248,
+      "eval_unseen_axis_steps_per_second": 156.195,
+      "epoch": 1.2360939431396787,
+      "step": 8000
+    },
+    {
+      "eval_hard_loss": 0.34612923860549927,
+      "eval_hard_accuracy": 0.8957983193277311,
+      "eval_hard_runtime": 0.2561,
+      "eval_hard_samples_per_second": 4646.719,
+      "eval_hard_steps_per_second": 74.191,
+      "epoch": 1.2360939431396787,
+      "step": 8000
+    },
+    {
+      "loss": 0.13536205291748046,
+      "grad_norm": 1.5142947435379028,
+      "learning_rate": 9.95121945299005e-05,
+      "epoch": 1.2669962917181705,
+      "step": 8200
+    },
+    {
+      "loss": 0.12326237678527832,
+      "grad_norm": 5.0126051902771,
+      "learning_rate": 9.227498688962761e-05,
+      "epoch": 1.2978986402966626,
+      "step": 8400
+    },
+    {
+      "loss": 0.12851189613342284,
+      "grad_norm": 3.0441036224365234,
+      "learning_rate": 8.519227031772971e-05,
+      "epoch": 1.3288009888751544,
+      "step": 8600
+    },
+    {
+      "loss": 0.11189216613769531,
+      "grad_norm": 1.7037580013275146,
+      "learning_rate": 7.828300048682143e-05,
+      "epoch": 1.3597033374536465,
+      "step": 8800
+    },
+    {
+      "loss": 0.12371747016906738,
+      "grad_norm": 0.7055680751800537,
+      "learning_rate": 7.156566886914933e-05,
+      "epoch": 1.3906056860321385,
+      "step": 9000
+    },
+    {
+      "eval_unseen_lanes_loss": 0.23261435329914093,
+      "eval_unseen_lanes_accuracy": 0.9221222463543283,
+      "eval_unseen_lanes_runtime": 0.5543,
+      "eval_unseen_lanes_samples_per_second": 5814.13,
+      "eval_unseen_lanes_steps_per_second": 92.001,
+      "epoch": 1.3906056860321385,
+      "step": 9000
+    },
+    {
+      "eval_unseen_domain_loss": 0.9834672808647156,
+      "eval_unseen_domain_accuracy": 0.6840893161109443,
+      "eval_unseen_domain_runtime": 1.579,
+      "eval_unseen_domain_samples_per_second": 6324.919,
+      "eval_unseen_domain_steps_per_second": 99.43,
+      "epoch": 1.3906056860321385,
+      "step": 9000
+    },
+    {
+      "eval_unseen_axis_loss": 1.5792756080627441,
+      "eval_unseen_axis_accuracy": 0.593996459037502,
+      "eval_unseen_axis_runtime": 1.2867,
+      "eval_unseen_axis_samples_per_second": 9657.248,
+      "eval_unseen_axis_steps_per_second": 151.55,
+      "epoch": 1.3906056860321385,
+      "step": 9000
+    },
+    {
+      "eval_hard_loss": 0.31886979937553406,
+      "eval_hard_accuracy": 0.9,
+      "eval_hard_runtime": 0.2641,
+      "eval_hard_samples_per_second": 4505.693,
+      "eval_hard_steps_per_second": 71.94,
+      "epoch": 1.3906056860321385,
+      "step": 9000
+    },
+    {
+      "loss": 0.11327756881713867,
+      "grad_norm": 2.5306780338287354,
+      "learning_rate": 6.505825324734749e-05,
+      "epoch": 1.4215080346106304,
+      "step": 9200
+    },
+    {
+      "loss": 0.11213048934936523,
+      "grad_norm": 3.7060656547546387,
+      "learning_rate": 5.8778169599994366e-05,
+      "epoch": 1.4524103831891224,
+      "step": 9400
+    },
+    {
+      "loss": 0.11299308776855468,
+      "grad_norm": 2.2346725463867188,
+      "learning_rate": 5.274222549074192e-05,
+      "epoch": 1.4833127317676142,
+      "step": 9600
+    },
+    {
+      "loss": 0.10625232696533203,
+      "grad_norm": 2.6063334941864014,
+      "learning_rate": 4.696657508576255e-05,
+      "epoch": 1.5142150803461063,
+      "step": 9800
+    },
+    {
+      "loss": 0.10512189865112305,
+      "grad_norm": 4.146987438201904,
+      "learning_rate": 4.146667591990211e-05,
+      "epoch": 1.5451174289245984,
+      "step": 10000
+    },
+    {
+      "eval_unseen_lanes_loss": 0.2341764271259308,
+      "eval_unseen_lanes_accuracy": 0.9286379149860379,
+      "eval_unseen_lanes_runtime": 0.5514,
+      "eval_unseen_lanes_samples_per_second": 5845.486,
+      "eval_unseen_lanes_steps_per_second": 92.498,
+      "epoch": 1.5451174289245984,
+      "step": 10000
+    },
+    {
+      "eval_unseen_domain_loss": 1.1408356428146362,
+      "eval_unseen_domain_accuracy": 0.6845899669570441,
+      "eval_unseen_domain_runtime": 1.5844,
+      "eval_unseen_domain_samples_per_second": 6303.222,
+      "eval_unseen_domain_steps_per_second": 99.089,
+      "epoch": 1.5451174289245984,
+      "step": 10000
+    },
+    {
+      "eval_unseen_axis_loss": 1.6913553476333618,
+      "eval_unseen_axis_accuracy": 0.6064703042008691,
+      "eval_unseen_axis_runtime": 1.2705,
+      "eval_unseen_axis_samples_per_second": 9780.492,
+      "eval_unseen_axis_steps_per_second": 153.484,
+      "epoch": 1.5451174289245984,
+      "step": 10000
+    },
+    {
+      "eval_hard_loss": 0.3436157703399658,
+      "eval_hard_accuracy": 0.9042016806722689,
+      "eval_hard_runtime": 0.2553,
+      "eval_hard_samples_per_second": 4661.691,
+      "eval_hard_steps_per_second": 74.43,
+      "epoch": 1.5451174289245984,
+      "step": 10000
+    },
+    {
+      "loss": 0.10526069641113281,
+      "grad_norm": 1.0956754684448242,
+      "learning_rate": 3.625724752724645e-05,
+      "epoch": 1.5760197775030902,
+      "step": 10200
+    },
+    {
+      "loss": 0.0974786376953125,
+      "grad_norm": 1.1272306442260742,
+      "learning_rate": 3.135223204682058e-05,
+      "epoch": 1.6069221260815822,
+      "step": 10400
+    },
+    {
+      "loss": 0.101556396484375,
+      "grad_norm": 3.4264066219329834,
+      "learning_rate": 2.6764756908851686e-05,
+      "epoch": 1.637824474660074,
+      "step": 10600
+    },
+    {
+      "loss": 0.10870396614074707,
+      "grad_norm": 1.7836824655532837,
+      "learning_rate": 2.2507099701460036e-05,
+      "epoch": 1.6687268232385661,
+      "step": 10800
+    },
+    {
+      "loss": 0.09275879859924316,
+      "grad_norm": 1.3714878559112549,
+      "learning_rate": 1.8590655311806385e-05,
+      "epoch": 1.6996291718170582,
+      "step": 11000
+    },
+    {
+      "eval_unseen_lanes_loss": 0.21319718658924103,
+      "eval_unseen_lanes_accuracy": 0.9304995345950977,
+      "eval_unseen_lanes_runtime": 0.5553,
+      "eval_unseen_lanes_samples_per_second": 5804.033,
+      "eval_unseen_lanes_steps_per_second": 91.842,
+      "epoch": 1.6996291718170582,
+      "step": 11000
+    },
+    {
+      "eval_unseen_domain_loss": 1.022688627243042,
+      "eval_unseen_domain_accuracy": 0.6976068889556424,
+      "eval_unseen_domain_runtime": 1.5825,
+      "eval_unseen_domain_samples_per_second": 6310.792,
+      "eval_unseen_domain_steps_per_second": 99.208,
+      "epoch": 1.6996291718170582,
+      "step": 11000
+    },
+    {
+      "eval_unseen_axis_loss": 1.5096590518951416,
+      "eval_unseen_axis_accuracy": 0.6129888942539836,
+      "eval_unseen_axis_runtime": 1.2878,
+      "eval_unseen_axis_samples_per_second": 9649.329,
+      "eval_unseen_axis_steps_per_second": 151.426,
+      "epoch": 1.6996291718170582,
+      "step": 11000
+    },
+    {
+      "eval_hard_loss": 0.31117716431617737,
+      "eval_hard_accuracy": 0.9008403361344538,
+      "eval_hard_runtime": 0.2589,
+      "eval_hard_samples_per_second": 4596.879,
+      "eval_hard_steps_per_second": 73.396,
+      "epoch": 1.6996291718170582,
+      "step": 11000
+    },
+    {
+      "loss": 0.09593048095703124,
+      "grad_norm": 3.177462100982666,
+      "learning_rate": 1.5025905429636426e-05,
+      "epoch": 1.73053152039555,
+      "step": 11200
+    },
+    {
+      "loss": 0.09772712707519532,
+      "grad_norm": 2.102013349533081,
+      "learning_rate": 1.1822390494841139e-05,
+      "epoch": 1.7614338689740419,
+      "step": 11400
+    },
+    {
+      "loss": 0.08881115913391113,
+      "grad_norm": 0.3349379897117615,
+      "learning_rate": 8.98868416411046e-06,
+      "epoch": 1.792336217552534,
+      "step": 11600
+    },
+    {
+      "loss": 0.09125228881835938,
+      "grad_norm": 1.5826376676559448,
+      "learning_rate": 6.532370365015982e-06,
+      "epoch": 1.823238566131026,
+      "step": 11800
+    },
+    {
+      "loss": 0.1032994556427002,
+      "grad_norm": 2.430258274078369,
+      "learning_rate": 4.460022998933166e-06,
+      "epoch": 1.854140914709518,
+      "step": 12000
+    },
+    {
+      "eval_unseen_lanes_loss": 0.21029867231845856,
+      "eval_unseen_lanes_accuracy": 0.9329816940738442,
+      "eval_unseen_lanes_runtime": 0.5539,
+      "eval_unseen_lanes_samples_per_second": 5818.497,
+      "eval_unseen_lanes_steps_per_second": 92.071,
+      "epoch": 1.854140914709518,
+      "step": 12000
+    },
+    {
+      "eval_unseen_domain_loss": 1.0129146575927734,
+      "eval_unseen_domain_accuracy": 0.700010013016922,
+      "eval_unseen_domain_runtime": 1.5764,
+      "eval_unseen_domain_samples_per_second": 6335.382,
+      "eval_unseen_domain_steps_per_second": 99.595,
+      "epoch": 1.854140914709518,
+      "step": 12000
+    },
+    {
+      "eval_unseen_axis_loss": 1.5864043235778809,
+      "eval_unseen_axis_accuracy": 0.6117817479478512,
+      "eval_unseen_axis_runtime": 1.2684,
+      "eval_unseen_axis_samples_per_second": 9796.563,
+      "eval_unseen_axis_steps_per_second": 153.737,
+      "epoch": 1.854140914709518,
+      "step": 12000
+    },
+    {
+      "eval_hard_loss": 0.30691978335380554,
+      "eval_hard_accuracy": 0.9084033613445378,
+      "eval_hard_runtime": 0.2602,
+      "eval_hard_samples_per_second": 4573.644,
+      "eval_hard_steps_per_second": 73.025,
+      "epoch": 1.854140914709518,
+      "step": 12000
+    },
+    {
+      "loss": 0.08737665176391601,
+      "grad_norm": 1.072799801826477,
+      "learning_rate": 2.7771883471252034e-06,
+      "epoch": 1.8850432632880099,
+      "step": 12200
+    },
+    {
+      "loss": 0.08102713584899902,
+      "grad_norm": 1.5238364934921265,
+      "learning_rate": 1.4883702270751197e-06,
+      "epoch": 1.9159456118665017,
+      "step": 12400
+    },
+    {
+      "loss": 0.09933463096618653,
+      "grad_norm": 0.16273926198482513,
+      "learning_rate": 5.970179387931151e-07,
+      "epoch": 1.9468479604449938,
+      "step": 12600
+    },
+    {
+      "loss": 0.10335576057434082,
+      "grad_norm": 0.17151758074760437,
+      "learning_rate": 1.0551703335823225e-07,
+      "epoch": 1.9777503090234858,
+      "step": 12800
+    },
+    {
+      "eval_unseen_lanes_loss": 0.21011516451835632,
+      "eval_unseen_lanes_accuracy": 0.9336022339435309,
+      "eval_unseen_lanes_runtime": 0.5534,
+      "eval_unseen_lanes_samples_per_second": 5823.753,
+      "eval_unseen_lanes_steps_per_second": 92.154,
+      "epoch": 2.0,
+      "step": 12944
+    },
+    {
+      "eval_unseen_domain_loss": 1.0266432762145996,
+      "eval_unseen_domain_accuracy": 0.7009111845399019,
+      "eval_unseen_domain_runtime": 1.5817,
+      "eval_unseen_domain_samples_per_second": 6314.155,
+      "eval_unseen_domain_steps_per_second": 99.261,
+      "epoch": 2.0,
+      "step": 12944
+    },
+    {
+      "eval_unseen_axis_loss": 1.6092302799224854,
+      "eval_unseen_axis_accuracy": 0.6111379365845807,
+      "eval_unseen_axis_runtime": 1.2767,
+      "eval_unseen_axis_samples_per_second": 9732.892,
+      "eval_unseen_axis_steps_per_second": 152.737,
+      "epoch": 2.0,
+      "step": 12944
+    },
+    {
+      "eval_hard_loss": 0.3078996539115906,
+      "eval_hard_accuracy": 0.9084033613445378,
+      "eval_hard_runtime": 0.2718,
+      "eval_hard_samples_per_second": 4377.996,
+      "eval_hard_steps_per_second": 69.901,
+      "epoch": 2.0,
+      "step": 12944
+    },
+    {
+      "train_runtime": 262.1825,
+      "train_samples_per_second": 1579.793,
+      "train_steps_per_second": 49.37,
+      "total_flos": 0.0,
+      "train_loss": 0.2282541081106707,
+      "epoch": 2.0,
+      "step": 12944
+    }
+  ]
+}

train_router_head.py

@@ -0,0 +1,218 @@
+"""
+Arabic zero-shot router with a real routing head.
+
+Two earlier formulations, and why this one:
+
+  fixed-slot head  - `num_labels = max_lanes` over one pooled vector. Scored exactly 1/n at every
+                     lane count. The head is positional (slot i's logit is w_i . h) and the corpus
+                     randomises lane order, so there is nothing to learn.
+  pairwise         - score each (text, category) separately, argmax. Works (0.842 on unseen lane
+                     sets) but costs n forward passes and each category is scored in isolation,
+                     never against its competitors.
+
+Here the categories live in the sequence and get their *own* pooled vectors, which a shared linear
+head turns into one logit each; the softmax is over the categories present. The head is shared
+across positions, so it scores category *content*, not slot index - which is what makes it
+position-invariant and zero-shot.
+
+Layout is text-first, categories-second, deliberately. The reference encoder is bidirectional so
+order does not matter there; our base is causal, and this ordering is what lets every category
+token attend to the whole text. Reversed, the categories would be encoded blind to the text.
+
+  النص:
+  {text}
+
+  الفئات:
+  - فئة أولى
+  - فئة ثانية
+
+Each category's character span is mapped to token indices via the tokenizer's offset mapping and
+mean-pooled. Slots beyond a row's category count are masked to -inf before the loss.
+"""
+import json
+import os
+from collections import defaultdict
+from pathlib import Path
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import Dataset
+from transformers import AutoModel, AutoTokenizer, Trainer, TrainingArguments
+
+os.environ.setdefault("CUDA_VISIBLE_DEVICES", "0")
+
+BASE_MODEL = os.environ.get("BASE_MODEL", "oddadmix/50M-2048-Emhotob")
+OUTPUT_DIR = os.environ.get("OUTPUT_DIR", "./Nawah-Router-head-50M")
+DATA = Path(os.environ.get("DATA_DIR", "data"))
+MAX_ROUTES = int(os.environ.get("MAX_ROUTES", 9))
+MAX_LENGTH = int(os.environ.get("MAX_LENGTH", 320))
+LR = float(os.environ.get("LR", 3e-4))
+EPOCHS = float(os.environ.get("EPOCHS", 3))
+BATCH_SIZE = int(os.environ.get("BATCH_SIZE", 32))
+WARMUP = int(os.environ.get("WARMUP", 500))
+SEED = 42
+
+
+def build_text(row):
+    """-> (full string, [(start, end) char span per category])."""
+    head = f"النص:\n{row['text']}\n\nالفئات:\n"
+    s = head
+    spans = []
+    for c in row["routes"]:
+        s += "- "
+        spans.append((len(s), len(s) + len(c)))
+        s += c + "\n"
+    return s, spans
+
+
+class RouterModel(nn.Module):
+    """Backbone + a shared linear scorer applied to each category's pooled span."""
+
+    def __init__(self, base_model):
+        super().__init__()
+        self.backbone = AutoModel.from_pretrained(base_model, dtype=torch.float32)
+        h = self.backbone.config.hidden_size
+        self.score = nn.Sequential(nn.Linear(h, h), nn.GELU(), nn.Linear(h, 1))
+        self.config = self.backbone.config
+
+    def forward(self, input_ids, attention_mask, cat_pool, n_routes, labels=None):
+        hs = self.backbone(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state
+        # cat_pool: (B, K, T) row-normalised selector -> (B, K, H)
+        cat_vecs = torch.bmm(cat_pool.to(hs.dtype), hs)
+        logits = self.score(cat_vecs).squeeze(-1)                       # (B, K)
+        ar = torch.arange(logits.size(1), device=logits.device)[None, :]
+        logits = logits.masked_fill(ar >= n_routes[:, None], torch.finfo(logits.dtype).min)
+        loss = F.cross_entropy(logits, labels) if labels is not None else None
+        return {"loss": loss, "logits": logits}
+
+
+def load(name):
+    p = DATA / f"{name}.jsonl"
+    return [json.loads(l) for l in open(p, encoding="utf-8")] if p.exists() else []
+
+
+class RouterDataset(Dataset):
+    def __init__(self, rows, tok, max_length):
+        self.rows, self.tok, self.max_length = rows, tok, max_length
+
+    def __len__(self):
+        return len(self.rows)
+
+    def __getitem__(self, i):
+        r = self.rows[i]
+        text, spans = build_text(r)
+        enc = self.tok(text, return_offsets_mapping=True, add_special_tokens=False,
+                       truncation=True, max_length=self.max_length)
+        ids, offs = enc["input_ids"], enc["offset_mapping"]
+        pool = torch.zeros(MAX_ROUTES, len(ids))
+        for ci, (s, e) in enumerate(spans[:MAX_ROUTES]):
+            idxs = [t for t, (a, b) in enumerate(offs) if a < e and b > s and a != b]
+            if idxs:
+                pool[ci, idxs] = 1.0 / len(idxs)
+        return {"input_ids": torch.tensor(ids, dtype=torch.long), "cat_pool": pool,
+                "n_routes": torch.tensor(min(len(r["routes"]), MAX_ROUTES), dtype=torch.long),
+                "labels": torch.tensor(min(r["label"], MAX_ROUTES - 1), dtype=torch.long)}
+
+
+class Collator:
+    def __init__(self, pad_id):
+        self.pad_id = pad_id
+
+    def __call__(self, feats):
+        T = max(f["input_ids"].size(0) for f in feats)
+        ids, att, pools = [], [], []
+        for f in feats:
+            n = f["input_ids"].size(0); pad = T - n
+            ids.append(torch.cat([f["input_ids"], torch.full((pad,), self.pad_id, dtype=torch.long)]))
+            att.append(torch.cat([torch.ones(n, dtype=torch.long), torch.zeros(pad, dtype=torch.long)]))
+            pools.append(F.pad(f["cat_pool"], (0, pad)))
+        return {"input_ids": torch.stack(ids), "attention_mask": torch.stack(att),
+                "cat_pool": torch.stack(pools),
+                "n_routes": torch.stack([f["n_routes"] for f in feats]),
+                "labels": torch.stack([f["labels"] for f in feats])}
+
+
+def metrics_fn(p):
+    return {"accuracy": float((np.asarray(p.predictions).argmax(-1) ==
+                               np.asarray(p.label_ids)).mean())}
+
+
+@torch.no_grad()
+def report(model, tok, rows, name, batch=64):
+    model.eval()
+    ds = RouterDataset(rows, tok, MAX_LENGTH); coll = Collator(tok.pad_token_id)
+    preds = []
+    for i in range(0, len(rows), batch):
+        b = coll([ds[j] for j in range(i, min(i + batch, len(rows)))])
+        b = {k: v.to(next(model.parameters()).device) for k, v in b.items()}
+        b.pop("labels")
+        preds += model(**b)["logits"].argmax(-1).tolist()
+    correct = [int(p == r["label"]) for p, r in zip(preds, rows)]
+    acc = sum(correct) / len(rows)
+    rand = sum(1 / len(r["routes"]) for r in rows) / len(rows)
+    print(f"\n[{name}]  n={len(rows):,}  route accuracy {acc:.4f}   random {rand:.4f}")
+    out = {"accuracy": acc, "random_baseline": rand, "n": len(rows)}
+    for key in ("n_routes", "difficulty", "mode"):
+        if key == "mode" and not all("mode" in r for r in rows):
+            continue
+        b = defaultdict(lambda: [0, 0])
+        for c, r in zip(correct, rows):
+            k = len(r["routes"]) if key == "n_routes" else r[key]
+            b[k][1] += 1; b[k][0] += c
+        print(f"    by {key:<10} " +
+              "  ".join(f"{k}:{v[0]/v[1]:.3f}(n={v[1]})" for k, v in sorted(b.items(), key=str)))
+        out[f"by_{key}"] = {str(k): {"acc": v[0]/v[1], "n": v[1]} for k, v in b.items()}
+    return out
+
+
+def main():
+    tok = AutoTokenizer.from_pretrained(BASE_MODEL)
+    if tok.pad_token_id is None:
+        tok.pad_token = tok.eos_token
+    model = RouterModel(BASE_MODEL)
+    print(f"[*] {BASE_MODEL} | params {sum(p.numel() for p in model.parameters())/1e6:.2f}M "
+          f"| hidden {model.config.hidden_size}")
+
+    train = load("train")
+    evals = {"unseen_lanes": load("eval_unseen_lanes"),
+             "unseen_domain": load("eval_unseen_domain"),
+             # v2 only: axes held out of training entirely. The strongest zero-shot test, so the
+             # checkpoint is selected on it when it exists.
+             "unseen_axis": load("eval_unseen_axis"),
+             "hard": load("eval_hard")}
+    evals = {k: v for k, v in evals.items() if v}
+    print(f"[*] train {len(train):,} | " + " | ".join(f"{k} {len(v):,}" for k, v in evals.items()))
+
+    args = TrainingArguments(
+        output_dir=OUTPUT_DIR, num_train_epochs=EPOCHS,
+        per_device_train_batch_size=BATCH_SIZE, per_device_eval_batch_size=64,
+        learning_rate=LR, lr_scheduler_type="cosine", warmup_steps=WARMUP,
+        max_grad_norm=1.0, bf16=True, logging_steps=200,
+        eval_strategy="steps", eval_steps=1000, save_strategy="steps", save_steps=1000,
+        save_total_limit=2, load_best_model_at_end=True,
+        metric_for_best_model=("eval_unseen_axis_accuracy" if "unseen_axis" in evals
+                               else "eval_unseen_domain_accuracy"), greater_is_better=True,
+        report_to=[], seed=SEED, dataloader_num_workers=4, remove_unused_columns=False,
+        label_names=["labels"])
+
+    trainer = Trainer(model=model, args=args,
+                      train_dataset=RouterDataset(train, tok, MAX_LENGTH),
+                      eval_dataset={k: RouterDataset(v, tok, MAX_LENGTH)
+                                    for k, v in evals.items() if v},
+                      data_collator=Collator(tok.pad_token_id), compute_metrics=metrics_fn)
+    trainer.train()
+
+    Path(OUTPUT_DIR).mkdir(exist_ok=True)
+    torch.save(model.state_dict(), Path(OUTPUT_DIR, "router_model.pt"))
+    model.backbone.save_pretrained(OUTPUT_DIR); tok.save_pretrained(OUTPUT_DIR)
+    results = {k: report(model, tok, v, k) for k, v in evals.items() if v}
+    Path(OUTPUT_DIR, "train_metrics.json").write_text(json.dumps(
+        {"results": results, "base_model": BASE_MODEL, "formulation": "routing_head",
+         "log_history": trainer.state.log_history}, ensure_ascii=False, indent=2), encoding="utf-8")
+    print(f"\n[+] done -> {OUTPUT_DIR}")
+
+
+if __name__ == "__main__":
+    main()