build_router_dataset_v2.py

"""
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()