Upload training/router_integration_benchmark.py with huggingface_hub

training/router_integration_benchmark.py

@@ -0,0 +1,237 @@
+#!/usr/bin/env python3
+"""
+Standalone integration: Replace heuristic router with trained XGBoost router.
+
+This script:
+1. Generates 2K eval traces (same as standalone_eval_v2.py)
+2. Runs the full benchmark with the TRAINED router replacing _route_learned()
+3. Compares: heuristic, trained, and oracle routers
+"""
+
+import json, os, sys, random, uuid, pickle
+import numpy as np
+from datetime import datetime, timedelta
+from dataclasses import dataclass, field
+from enum import Enum
+from collections import defaultdict
+from typing import Dict, List, Optional, Any
+
+# ─── Load trained models ─────────────────────────────────────────────
+print("="*80)
+print("ACO TRAINED ROUTER - INTEGRATION BENCHMARK")
+print("="*80)
+
+import xgboost as xgb
+
+MODEL_DIR = "/app/router_models"
+feat_keys = json.load(open(f"{MODEL_DIR}/feat_keys.json"))
+tier_config = json.load(open(f"{MODEL_DIR}/tier_config.json"))
+TIER_COST = {int(k):v for k,v in tier_config["tier_cost"].items()}
+TIER_STR = {int(k):v for k,v in tier_config["tier_str"].items()}
+TASK_FLOOR = tier_config["task_floor"]
+
+print(f"\n[1] Loading trained router models...")
+tier_clfs = {}
+for tier in range(1, 6):
+    clf = xgb.XGBClassifier()
+    clf.load_model(f"{MODEL_DIR}/tier_{tier}_success.json")
+    tier_clfs[tier] = clf
+    print(f"  Loaded tier_{tier}_success.json")
+
+# ─── Feature extraction (must match training) ────────────────────────
+TASK_TYPES = ["quick_answer","coding","research","document_drafting",
+              "legal_regulated","tool_heavy","retrieval_heavy",
+              "long_horizon","unknown_ambiguous"]
+TT2IDX = {t:i for i,t in enumerate(TASK_TYPES)}
+
+CODE_KW = ["python","javascript","code","function","bug","debug","refactor",
+           "implement","test","compile","runtime","class","module","async","thread"]
+LEGAL_KW = ["contract","legal","compliance","gdpr","privacy","policy","regulatory","liability"]
+RESEARCH_KW = ["research","find sources","literature","investigate","compare","analyze","survey"]
+TOOL_KW = ["search","fetch","retrieve","query","api","database","scrape","aggregate"]
+LONG_KW = ["plan","project","roadmap","orchestrate","multi-step","migrate","pipeline","deploy"]
+MATH_KW = ["calculate","compute","solve","equation","formula","optimize","probability"]
+
+def extract_features(request, task_type, difficulty=3):
+    r = request.lower()
+    f = {"req_len": len(request), "num_words": len(request.split()),
+        "has_code": int(any(k in r for k in CODE_KW)),
+        "n_code": sum(1 for k in CODE_KW if k in r),
+        "has_legal": int(any(k in r for k in LEGAL_KW)),
+        "n_legal": sum(1 for k in LEGAL_KW if k in r),
+        "has_research": int(any(k in r for k in RESEARCH_KW)),
+        "n_research": sum(1 for k in RESEARCH_KW if k in r),
+        "has_tool": int(any(k in r for k in TOOL_KW)),
+        "n_tool": sum(1 for k in TOOL_KW if k in r),
+        "has_long": int(any(k in r for k in LONG_KW)),
+        "has_math": int(any(k in r for k in MATH_KW)),
+        "tt_idx": TT2IDX.get(task_type, 8), "difficulty": difficulty}
+    for tt in TASK_TYPES:
+        f[f"tt_{tt}"] = int(task_type == tt)
+    return f
+
+def f2v(feats):
+    return np.array([float(feats.get(k, 0.0)) for k in feat_keys], dtype=np.float32)
+
+# ─── Routing Functions ────────────────────────────────────────────────
+def route_trained(request, task_type, difficulty):
+    """Trained router: per-tier P(success) + safety floor + asymmetric cost."""
+    feats = extract_features(request, task_type, difficulty)
+    x = f2v(feats).reshape(1, -1)
+    floor = TASK_FLOOR.get(task_type, 2)
+    
+    best_tier = floor; best_score = float("inf")
+    for tier in range(floor, 6):
+        p_success = tier_clfs[tier].predict_proba(x)[0, 1]
+        p_fail = 1.0 - p_success
+        cost_norm = TIER_COST[tier] / TIER_COST[5]
+        score = p_fail * 5.0 + cost_norm * 1.0  # asymmetric: 5x underkill penalty
+        if score < best_score:
+            best_score = score; best_tier = tier
+    return best_tier
+
+def route_heuristic(task_type, difficulty):
+    """Original heuristic router: difficulty + 1."""
+    return min(difficulty + 1, 5)
+
+def route_frontier():
+    return 4
+
+def route_cascade_trained(request, task_type, difficulty):
+    """Cascade: start at floor, escalate if P(success) < threshold."""
+    feats = extract_features(request, task_type, difficulty)
+    x = f2v(feats).reshape(1, -1)
+    floor = TASK_FLOOR.get(task_type, 2)
+    
+    for tier in range(floor, 6):
+        p_success = tier_clfs[tier].predict_proba(x)[0, 1]
+        if p_success >= 0.65:
+            return tier
+    return 4  # fallback to frontier
+
+# ─── Generate Evaluation Traces ────────────────────────────────────────
+print("\n[2] Generating 2K evaluation traces (different seed from training)...")
+
+class TaskType(Enum):
+    QUICK_ANSWER="quick_answer"; CODING="coding"; RESEARCH="research"
+    DOCUMENT_DRAFTING="document_drafting"; LEGAL_REGULATED="legal_regulated"
+    TOOL_HEAVY="tool_heavy"; RETRIEVAL_HEAVY="retrieval_heavy"
+    LONG_HORIZON="long_horizon"; UNKNOWN_AMBIGUOUS="unknown_ambiguous"
+
+TASK_TEMPLATES_EVAL = {
+    "quick_answer":["What is the capital of France?","Explain quantum computing briefly.","What is 237*452?"],
+    "coding":["Write a Python function to reverse a linked list.","Fix the bug in this React component.",
+              "Implement LRU cache in Go.","Debug segfault in C++ thread pool."],
+    "research":["Research latest transformer advances.","Find sources comparing LoRA and full FT.",
+                "Investigate data center climate impact."],
+    "document_drafting":["Draft project proposal for ML pipeline.","Write email to team about deployment."],
+    "legal_regulated":["Review this contract for liability clauses.","Check GDPR compliance for data pipeline.",
+                       "Draft privacy policy section."],
+    "tool_heavy":["Search open issues and create summary.","Fetch API docs and generate client code."],
+    "retrieval_heavy":["Answer based on 50-page document.","Find all payment processing mentions."],
+    "long_horizon":["Plan 3-month roadmap.","Orchestrate multi-region deployment."],
+    "unknown_ambiguous":["Help me with this thing.","I need something about the server."],
+}
+
+def tsp(tier, diff):
+    s = {1:0.35,2:0.55,3:0.80,4:0.93,5:0.97}[tier]
+    return s ** (diff * 0.6)
+
+eval_rng = random.Random(999)  # DIFFERENT seed for eval
+eval_traces = []
+for i in range(2000):
+    tt = eval_rng.choice(list(TASK_TEMPLATES_EVAL.keys()))
+    diff = {"quick_answer":1,"document_drafting":2,"tool_heavy":2,"retrieval_heavy":2,
+            "research":3,"coding":3,"unknown_ambiguous":3,"long_horizon":4,"legal_regulated":5}[tt]
+    tier_out = {t: eval_rng.random() < tsp(t, diff) for t in range(1,6)}
+    opt = 5
+    for t in range(1,6):
+        if tier_out[t]: opt = t; break
+    req = eval_rng.choice(TASK_TEMPLATES_EVAL[tt])
+    eval_traces.append({"tt":tt,"diff":diff,"opt":opt,"tier_out":tier_out,"req":req})
+
+print(f"  Generated {len(eval_traces)} eval traces")
+
+# ─── Evaluate All Routers ─────────────────────────────────────────────
+print("\n[3] Evaluating all routers on 2K traces...")
+
+def eval_router(name, route_fn):
+    succ = 0; cost = 0.0; unsafe = 0; fd = 0
+    td = defaultdict(int)
+    for t in eval_traces:
+        pred = route_fn(t)
+        td[pred] += 1
+        if t["tier_out"].get(pred, False):
+            succ += 1
+        elif pred < t["opt"]:
+            unsafe += 1
+        else:
+            fd += 1
+        cost += TIER_COST[pred]
+    n = len(eval_traces)
+    return {"success":succ/n, "avg_cost":cost/n, "unsafe_rate":unsafe/n,
+            "false_done":fd/n, "tier_dist":dict(td)}
+
+routers = {
+    "always_frontier": lambda t: 4,
+    "always_cheap": lambda t: 1,
+    "heuristic_diff+1": lambda t: min(t["diff"]+1, 5),
+    "heuristic_floor": lambda t: TASK_FLOOR.get(t["tt"], 2),
+    "trained_asymmetric": lambda t: route_trained(t["req"], t["tt"], t["diff"]),
+    "trained_cascade_t0.65": lambda t: route_cascade_trained(t["req"], t["tt"], t["diff"]),
+    "oracle": lambda t: t["opt"],
+}
+
+results = {}
+for name, fn in routers.items():
+    results[name] = eval_router(name, fn)
+    r = results[name]
+    fc = results["always_frontier"]["avg_cost"]
+    cr = (1 - r["avg_cost"]/fc)*100
+    print(f"  {name:<25} success={r['success']:.3f} cost={r['avg_cost']:.4f} costRed={cr:.1f}% unsafe={r['unsafe_rate']:.3f}")
+
+# ─── Per-Task-Type Breakdown ──────────────────────────────────────────
+print("\n\n[4] Per-task-type breakdown (trained vs heuristic vs frontier)...")
+for tt in sorted(set(t["tt"] for t in eval_traces)):
+    tt_traces = [t for t in eval_traces if t["tt"] == tt]
+    n_tt = len(tt_traces)
+    if n_tt == 0: continue
+    
+    for rname, rfn in [("frontier", lambda t:4),
+                       ("heuristic", lambda t:min(t["diff"]+1,5)),
+                       ("trained", lambda t:route_trained(t["req"],t["tt"],t["diff"]))]:
+        succ = sum(1 for t in tt_traces if t["tier_out"].get(rfn(t), False))
+        cost = sum(TIER_COST[rfn(t)] for t in tt_traces)
+        sr = succ/n_tt; ac = cost/n_tt
+        if rname == "frontier":
+            print(f"\n  {tt} (n={n_tt}):")
+        print(f"    {rname:<12} success={sr:.3f} cost={ac:.4f}")
+
+# ─── Final Comparison ─────────────────────────────────────────────────
+print(f"\n\n{'='*80}")
+print("FINAL INTEGRATION BENCHMARK")
+print(f"{'='*80}")
+print(f"\n{'Router':<25} {'Success':>10} {'AvgCost':>10} {'CostRed':>10} {'Unsafe':>10} {'F-DONE':>10}")
+print("-"*75)
+fc = results["always_frontier"]["avg_cost"]
+for name, r in sorted(results.items(), key=lambda x: (-x[1]["success"], x[1]["avg_cost"])):
+    cr = (1 - r["avg_cost"]/fc)*100
+    print(f"{name:<25} {r['success']:>10.3f} {r['avg_cost']:>10.4f} {cr:>9.1f}% {r['unsafe_rate']:>10.3f} {r['false_done']:>10.3f}")
+
+# Quality/cost frontier
+print("\nPARETO FRONTIER:")
+pareto = []
+for name, r in results.items():
+    if name == "always_cheap": continue
+    dominated = False
+    for name2, r2 in results.items():
+        if name == name2: continue
+        if r2["success"] >= r["success"] and r2["avg_cost"] <= r["avg_cost"]:
+            if r2["success"] > r["success"] or r2["avg_cost"] < r["avg_cost"]:
+                dominated = True; break
+    if not dominated:
+        pareto.append((name, r))
+        cr = (1 - r["avg_cost"]/fc)*100
+        print(f"  {name:<25} success={r['success']:.3f} cost={r['avg_cost']:.4f} costRed={cr:.1f}%")
+
+print(f"\n\nDONE!")