Upload training/router_v6_hybrid.py with huggingface_hub

training/router_v6_hybrid.py

@@ -0,0 +1,338 @@
+#!/usr/bin/env python3
+"""Trained Router v6: Hybrid heuristic + ML safety net.
+
+Key insight from v5: The ML classifiers alone can't beat the heuristic
+because difficulty is the dominant feature. But they CAN detect when
+the heuristic is wrong.
+
+Architecture:
+1. Heuristic: difficulty+1 with safety floor (this is the base)
+2. ML SAFETY NET: Check if P(success@heuristic_tier) < LOW_THRESHOLD
+   If so, escalate to next tier (the ML caught a case the heuristic missed)
+3. ML COST SAVER: Check if P(success@tier-1) >= HIGH_THRESHOLD
+   If so, DOWNGRADE one tier (the ML says a cheaper tier would work)
+"""
+import json, os, sys, random, uuid, pickle
+import numpy as np
+from collections import defaultdict
+from typing import Dict, List, Optional, Tuple
+
+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"]
+
+TIER_STR = {1:0.35,2:0.55,3:0.80,4:0.93,5:0.97}
+TIER_COST = {1:0.05,2:0.15,3:0.75,4:1.0,5:1.5}
+TASK_FLOOR = {"legal_regulated":4,"long_horizon":3,"research":3,"coding":3,
+              "unknown_ambiguous":3,"quick_answer":1,"document_drafting":2,
+              "tool_heavy":2,"retrieval_heavy":2}
+
+TASK_TEMPLATES = {
+    "quick_answer":["What is the capital of France?","Explain quantum computing briefly.",
+        "What is 237*452?","Define photosynthesis.","Who wrote Hamlet?",
+        "What is the speed of light?","List the primary colors.","What is GDP?"],
+    "coding":["Write a Python function to reverse a linked list.",
+        "Fix the bug in this React component.","Refactor auth module to JWT.",
+        "Implement LRU cache in Go.","Debug segfault in C++ thread pool.",
+        "Add unit tests for the payment module.","Optimize this SQL query.",
+        "Create a REST API for user management.","Implement binary search in Rust."],
+    "research":["Research latest transformer advances.",
+        "Find sources comparing LoRA and full FT.",
+        "Investigate data center climate impact.",
+        "Survey privacy-preserving ML techniques."],
+    "document_drafting":["Draft project proposal for ML pipeline.",
+        "Write email to team about deployment.","Create technical report on performance."],
+    "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.","Query Q3 sales and produce chart."],
+    "retrieval_heavy":["Answer based on 50-page document.",
+        "Find all payment processing mentions.","Retrieve relevant cases for legal query."],
+    "long_horizon":["Plan 3-month roadmap.","Orchestrate multi-region deployment.",
+        "Redesign data architecture end-to-end."],
+    "unknown_ambiguous":["Help me with this thing.",
+        "I need something about the server.","Can you look into that issue?"],
+}
+
+def tsp(tier, diff):
+    return TIER_STR[tier] ** (diff * 0.6)
+
+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 gen_trace(idx, rng):
+    tt = rng.choice(list(TASK_TEMPLATES.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: 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
+    if diff <= 2: actual = rng.choices([1,2,3,4,5],weights=[3,4,2,1,0.5])[0]
+    elif diff == 3: actual = rng.choices([1,2,3,4,5],weights=[1,2,4,2,1])[0]
+    elif diff == 4: actual = rng.choices([1,2,3,4,5],weights=[0.5,1,2,4,2])[0]
+    else: actual = rng.choices([1,2,3,4,5],weights=[0.2,0.5,1,3,4])[0]
+    outcome = "success" if tier_out[actual] else "failure"
+    req = rng.choice(TASK_TEMPLATES[tt])
+    feats = extract_features(req, tt, diff)
+    return {"feats":feats,"opt":opt,"actual":actual,"outcome":outcome,
+            "tier_out":tier_out,"tt":tt,"diff":diff,"req":req}
+
+print("="*80)
+print("ACO TRAINED ROUTER v6: HYBRID HEURISTIC + ML SAFETY NET")
+print("="*80)
+
+# ─── Train Models ────────────────────────────────────────────────────
+print("\n[1] Generating 50K training traces...")
+rng = random.Random(42)
+traces = [gen_trace(i, rng) for i in range(50000)]
+FEAT_KEYS = sorted(traces[0]["feats"].keys())
+def f2v(feats):
+    return np.array([float(feats.get(k, 0.0)) for k in FEAT_KEYS], dtype=np.float32)
+
+X_all = np.array([f2v(t["feats"]) for t in traces])
+y_opt = np.array([t["opt"] for t in traces])
+per_tier_labels = {}
+for tier in range(1,6):
+    per_tier_labels[tier] = np.array([1 if t["tier_out"].get(tier,False) else 0 for t in traces])
+
+from sklearn.model_selection import train_test_split
+from sklearn.metrics import accuracy_score, f1_score
+from sklearn.calibration import IsotonicRegression
+import xgboost as xgb
+
+X_train, X_test, idx_train, idx_test = train_test_split(X_all, range(len(traces)), test_size=0.2, random_state=42, stratify=y_opt)
+print(f"  Train: {len(X_train)}, Test: {len(X_test)}")
+
+print("\n[2] Training per-tier classifiers...")
+tier_clfs = {}
+tier_calibs = {}
+for tier in range(1,6):
+    y_tr = per_tier_labels[tier][idx_train]
+    y_te = per_tier_labels[tier][idx_test]
+    neg = (y_tr==0).sum(); pos = (y_tr==1).sum()
+    spw = neg/max(pos,1)
+    clf = xgb.XGBClassifier(n_estimators=200,max_depth=5,learning_rate=0.1,
+        subsample=0.8,colsample_bytree=0.8,scale_pos_weight=min(spw,5.0),
+        objective="binary:logistic",eval_metric="logloss",random_state=42,verbosity=0)
+    clf.fit(X_train, y_tr)
+    y_prob = clf.predict_proba(X_test)[:,1]
+    iso = IsotonicRegression(out_of_bounds="clip")
+    iso.fit(y_prob, y_te)
+    tier_clfs[tier] = clf
+    tier_calibs[tier] = iso
+    acc = accuracy_score(y_te, clf.predict(X_test))
+    f1 = f1_score(y_te, clf.predict(X_test), zero_division=0)
+    print(f"  Tier {tier}: acc={acc:.3f}, f1={f1:.3f}")
+
+def get_calibrated_psuccess(x, tier):
+    """Get calibrated P(success@tier) for a feature vector."""
+    p_raw = tier_clfs[tier].predict_proba(x)[0, 1]
+    return float(tier_calibs[tier].transform([p_raw])[0])
+
+# ─── Hybrid Router ────────────────────────────────────────────────────
+print("\n[3] Building hybrid heuristic + ML safety net router...")
+
+def route_hybrid(request, task_type, difficulty, 
+                 safety_threshold=0.35, downgrade_threshold=0.80):
+    """Hybrid: heuristic base + ML safety net + ML cost saver.
+    
+    1. Start with heuristic tier (difficulty+1, safety floor)
+    2. SAFETY NET: If P(success@heuristic_tier) < safety_threshold, ESCALATE
+    3. COST SAVER: If P(success@tier-1) >= downgrade_threshold AND
+       tier-1 >= safety_floor, DOWNGRADE one tier
+    4. Never go below safety floor or above 5
+    """
+    heuristic_tier = min(difficulty + 1, 5)
+    floor = TASK_FLOOR.get(task_type, 2)
+    heuristic_tier = max(heuristic_tier, floor)
+    
+    feats = extract_features(request, task_type, difficulty)
+    x = f2v(feats).reshape(1, -1)
+    
+    tier = heuristic_tier
+    
+    # SAFETY NET: Check if heuristic tier is likely to fail
+    p_success = get_calibrated_psuccess(x, tier)
+    if p_success < safety_threshold and tier < 5:
+        tier += 1
+        p_success = get_calibrated_psuccess(x, tier)
+    
+    # COST SAVER: Check if a cheaper tier would also work
+    # Only downgrade if: cheaper tier >= floor, P(success) high, and we're not already escalated
+    if tier > floor and tier == heuristic_tier:  # only if we didn't escalate
+        cheaper_tier = tier - 1
+        p_cheaper = get_calibrated_psuccess(x, cheaper_tier)
+        if p_cheaper >= downgrade_threshold and cheaper_tier >= floor:
+            tier = cheaper_tier
+    
+    return tier
+
+# ─── Generate Eval ─────────────────────────────────────────────────────
+print("\n[4] Generating 2K eval traces (seed=999)...")
+eval_rng = random.Random(999)
+eval_traces = []
+for i in range(2000):
+    tt = eval_rng.choice(list(TASK_TEMPLATES.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[tt])
+    eval_traces.append({"tt":tt,"diff":diff,"opt":opt,"tier_out":tier_out,"req":req})
+print(f"  Generated {len(eval_traces)} traces")
+
+# ─── Evaluate ──────────────────────────────────────────────────────────
+print("\n[5] Evaluating all routers...")
+n_test = len(eval_traces)
+results = {}
+
+def eval_router(name, route_fn):
+    succ=0; cost=0.0; unsafe=0; fd=0; td=defaultdict(int)
+    escalations=0; downgrades=0; heuristic_only=0
+    for t in eval_traces:
+        pred = route_fn(t)
+        td[pred] += 1
+        h_tier = min(t["diff"]+1, 5)
+        h_tier = max(h_tier, TASK_FLOOR.get(t["tt"], 2))
+        if pred > h_tier: escalations += 1
+        elif pred < h_tier: downgrades += 1
+        else: heuristic_only += 1
+        if t["tier_out"].get(pred, False): succ += 1
+        elif pred < t["opt"]: unsafe += 1
+        else: fd += 1
+        cost += TIER_COST[pred]
+    return {"success":succ/n_test, "avg_cost":cost/n_test, "unsafe_rate":unsafe/n_test,
+            "false_done":fd/n_test, "tier_dist":dict(td),
+            "escalations":escalations, "downgrades":downgrades, "heuristic_only":heuristic_only}
+
+results["always_frontier"] = eval_router("always_frontier", lambda t: 4)
+results["always_cheap"] = eval_router("always_cheap", lambda t: 1)
+results["heuristic_diff+1"] = eval_router("heuristic_diff+1", lambda t: min(t["diff"]+1, 5))
+results["heuristic_floor"] = eval_router("heuristic_floor", lambda t: TASK_FLOOR.get(t["tt"], 2))
+results["oracle"] = eval_router("oracle", lambda t: t["opt"])
+
+# Hybrid at different thresholds
+for st in [0.25, 0.30, 0.35, 0.40]:
+    for dt in [0.70, 0.75, 0.80, 0.85]:
+        name = f"hybrid_s{st:.2f}_d{dt:.2f}"
+        results[name] = eval_router(name,
+            lambda t, s=st, d=dt: route_hybrid(t["req"], t["tt"], t["diff"], s, d))
+
+# Print top results
+print(f"\n{'Router':<30} {'Success':>10} {'AvgCost':>10} {'CostRed':>10} {'Unsafe':>10} {'F-DONE':>10}")
+print("-"*80)
+fc = results["always_frontier"]["avg_cost"]
+
+# Only show key results + top 10 hybrids
+shown = set()
+for name in ["always_frontier","always_cheap","heuristic_diff+1","heuristic_floor","oracle"]:
+    r = results[name]
+    cr = (1 - r["avg_cost"]/fc)*100
+    print(f"{name:<30} {r['success']:>10.3f} {r['avg_cost']:>10.4f} {cr:>9.1f}% {r['unsafe_rate']:>10.3f} {r['false_done']:>10.3f}")
+    shown.add(name)
+
+# Top 10 hybrids by composite score
+hybrid_scores = []
+for name, r in results.items():
+    if name in shown or not name.startswith("hybrid"): continue
+    score = r["success"]*20 - r["avg_cost"]*30 - r["unsafe_rate"]*100
+    hybrid_scores.append((score, name, r))
+hybrid_scores.sort(reverse=True)
+
+for score, name, r in hybrid_scores[:10]:
+    cr = (1 - r["avg_cost"]/fc)*100
+    esc = r["escalations"]; down = r["downgrades"]; honly = r["heuristic_only"]
+    print(f"{name:<30} {r['success']:>10.3f} {r['avg_cost']:>10.4f} {cr:>9.1f}% {r['unsafe_rate']:>10.3f} {r['false_done']:>10.3f}  esc={esc} down={down} same={honly}")
+
+# ─── Per-task breakdown for best hybrid ────────────────────────────────
+best_hybrid_name = hybrid_scores[0][1] if hybrid_scores else "heuristic_diff+1"
+print(f"\n\n[6] Per-task breakdown for {best_hybrid_name}...")
+
+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
+    print(f"\n  {tt} (n={n_tt}):")
+    for rname, rfn in [("frontier", lambda t:4),
+                       ("heuristic", lambda t:min(t["diff"]+1,5)),
+                       ("hybrid", lambda t:route_hybrid(t["req"],t["tt"],t["diff"],0.35,0.80)),
+                       ("oracle", lambda t:t["opt"])]:
+        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
+        cr = (1-ac/fc)*100
+        print(f"    {rname:<12} success={sr:.3f} cost={ac:.4f} costRed={cr:.1f}%")
+
+# ─── Pareto ────────────────────────────────────────────────────────────
+print(f"\n\n[7] Pareto frontier...")
+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:
+        cr = (1-r["avg_cost"]/fc)*100
+        print(f"  {name:<30} success={r['success']:.3f} cost={r['avg_cost']:.4f} costRed={cr:.1f}% unsafe={r['unsafe_rate']:.3f}")
+
+# ─── Save ──────────────────────────────────────────────────────────────
+print("\n[8] Saving final model...")
+os.makedirs("/app/router_models", exist_ok=True)
+
+bundle = {
+    "tier_clfs": {str(k):v for k,v in tier_clfs.items()},
+    "tier_calibrators": {str(k):v for k,v in tier_calibs.items()},
+    "feat_keys": FEAT_KEYS,
+    "tier_config": {"tier_cost":TIER_COST,"tier_str":TIER_STR,
+                    "task_floor":TASK_FLOOR,
+                    "safety_threshold":0.35,"downgrade_threshold":0.80},
+    "version": "6.0",
+    "description": "ACO Hybrid Router: heuristic base + ML safety net + ML cost saver",
+}
+with open("/app/router_models/router_bundle_v6.pkl","wb") as f:
+    pickle.dump(bundle, f)
+print(f"  Saved router_bundle_v6.pkl ({os.path.getsize('/app/router_models/router_bundle_v6.pkl')/1024:.0f} KB)")
+
+with open("/app/router_models/v6_eval_results.json","w") as f:
+    json.dump(results, f, indent=2, default=str)
+
+print(f"\n\n{'='*80}")
+print("FINAL v6 COMPARISON")
+print(f"{'='*80}")
+print(f"\n{'Router':<30} {'Success':>10} {'AvgCost':>10} {'CostRed':>10} {'Unsafe':>10} {'F-DONE':>10}")
+print("-"*80)
+for name, r in sorted(results.items(), key=lambda x: (-x[1]["success"], x[1]["avg_cost"])):
+    if name.startswith("hybrid") and name != best_hybrid_name: continue
+    cr = (1-r["avg_cost"]/fc)*100
+    print(f"{name:<30} {r['success']:>10.3f} {r['avg_cost']:>10.4f} {cr:>9.1f}% {r['unsafe_rate']:>10.3f} {r['false_done']:>10.3f}")
+
+print(f"\nDONE!")