Upload training/train_router_v4_production.py with huggingface_hub

training/train_router_v4_production.py

@@ -0,0 +1,241 @@
+#!/usr/bin/env python3
+"""Production Trained Router v4: Difficulty-first + ML refinement.
+
+Architecture:
+1. Predict difficulty (task_type -> difficulty mapping + ML difficulty classifier)
+2. Convert difficulty to base_tier (difficulty + 1, capped at 5)
+3. Apply safety floor per task_type
+4. Use per-tier P(success) classifiers to CONFIRM or ESCALATE
+5. If P(success@base_tier) < 0.65, escalate to next tier
+
+This combines the best of both worlds:
+- Difficulty mapping (heuristic) is reliable and interpretable
+- ML classifiers add a safety net: they catch cases where difficulty is underestimated
+- Safety floors prevent dangerous under-routing on legal/critical tasks
+"""
+import json, os, sys, random, uuid
+import numpy as np
+from datetime import datetime
+from collections import defaultdict
+from typing import Dict, List, Optional, Any
+
+import xgboost as xgb
+
+# ─── Load Models ──────────────────────────────────────────────────────
+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"]
+
+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
+
+# ─── Feature Extraction ────────────────────────────────────────────────
+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)
+
+# ─── Production Router ─────────────────────────────────────────────────
+def route_production(request, task_type, difficulty, escalation_threshold=0.65):
+    """Production router: difficulty-first + ML confirmation.
+    
+    Step 1: difficulty -> base_tier (difficulty + 1, capped at 5)
+    Step 2: Apply safety floor (task_type -> minimum tier)
+    Step 3: base_tier = max(base_tier, safety_floor)
+    Step 4: If P(success@base_tier) < escalation_threshold, escalate one tier
+    Step 5: Never exceed tier 5
+    """
+    base_tier = min(difficulty + 1, 5)
+    floor = TASK_FLOOR.get(task_type, 2)
+    base_tier = max(base_tier, floor)
+    
+    # ML confirmation: check if base_tier is likely to succeed
+    feats = extract_features(request, task_type, difficulty)
+    x = f2v(feats).reshape(1, -1)
+    
+    p_success = tier_clfs[base_tier].predict_proba(x)[0, 1]
+    
+    # If P(success) at base_tier is too low, escalate
+    while p_success < escalation_threshold and base_tier < 5:
+        base_tier += 1
+        p_success = tier_clfs[base_tier].predict_proba(x)[0, 1]
+    
+    return base_tier
+
+# ─── Generate Evaluation ────────────────────────────────────────────────
+TASK_TEMPLATES_EVAL = {
+    "quick_answer":["What is the capital of France?","Explain quantum computing briefly.",
+        "What is 237*452?","Define photosynthesis.","Who wrote Hamlet?"],
+    "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."],
+    "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.","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):
+    s = {1:0.35,2:0.55,3:0.80,4:0.93,5:0.97}[tier]
+    return s ** (diff * 0.6)
+
+print("="*80)
+print("PRODUCTION TRAINED ROUTER v4 BENCHMARK")
+print("="*80)
+print("\nArchitecture: difficulty-first + ML confirmation + safety floors")
+
+print("\n[1] 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_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)} traces")
+
+# ─── Evaluate ──────────────────────────────────────────────────────────
+print("\n[2] Evaluating all routers...")
+
+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)}
+
+# Baselines
+results = {}
+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"])
+
+# Production router at different escalation thresholds
+for threshold in [0.50, 0.55, 0.60, 0.65, 0.70, 0.75]:
+    name = f"prod_t{threshold:.2f}"
+    results[name] = eval_router(name, 
+        lambda t, th=threshold: route_production(t["req"], t["tt"], t["diff"], escalation_threshold=th))
+
+# Print comparison
+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}")
+
+# Find best production config
+print("\n\n[3] Finding best production config...")
+best_name = ""
+best_score = -float("inf")
+for name, r in results.items():
+    if name in ("oracle","always_cheap"): continue
+    # Composite: maximize success, minimize cost and unsafe
+    score = r["success"]*20 - r["avg_cost"]*30 - r["unsafe_rate"]*100
+    if score > best_score:
+        best_score = score
+        best_name = name
+
+print(f"  Best: {best_name} (composite: {best_score:.2f})")
+print(f"  Success: {results[best_name]['success']:.3f}")
+print(f"  Avg cost: {results[best_name]['avg_cost']:.4f}")
+fc_val = results["always_frontier"]["avg_cost"]
+cr = (1 - results[best_name]["avg_cost"]/fc_val)*100
+print(f"  Cost reduction vs frontier: {cr:.1f}%")
+print(f"  Unsafe rate: {results[best_name]['unsafe_rate']:.3f}")
+print(f"  False-DONE rate: {results[best_name]['false_done']:.3f}")
+
+# Per-task breakdown for best
+print(f"\n\n[4] Per-task breakdown for {best_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
+    
+    for rname, rfn in [("frontier", lambda t:4),
+                       ("heuristic", lambda t:min(t["diff"]+1,5)),
+                       ("trained", lambda t:route_production(t["req"],t["tt"],t["diff"],
+                           escalation_threshold=float(best_name.split("t")[1]))),
+                       ("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
+        if rname == "frontier":
+            print(f"\n  {tt} (n={n_tt}):")
+        cr = (1 - ac/fc_val)*100
+        print(f"    {rname:<12} success={sr:.3f} cost={ac:.4f} costRed={cr:.1f}%")
+
+# Pareto
+print("\n\nPARETO 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_val)*100
+        print(f"  {name:<25} success={r['success']:.3f} cost={r['avg_cost']:.4f} costRed={cr:.1f}% unsafe={r['unsafe_rate']:.3f}")
+
+print(f"\nDONE!")