Upload training/router_v5_calibrated.py with huggingface_hub

training/router_v5_calibrated.py

@@ -0,0 +1,393 @@
+#!/usr/bin/env python3
+"""Trained Router v5: Calibrated + per-task thresholds + oversampled training.
+
+Key improvements over v4:
+1. Platt scaling calibration on held-out data
+2. Per-task-type escalation thresholds
+3. Oversampled easy-task successes in training data
+"""
+import json, os, sys, random, uuid
+import numpy as np
+from datetime import datetime
+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}
+
+# Per-task-type escalation thresholds (lower = more aggressive cost savings)
+TASK_THRESHOLDS = {
+    "quick_answer": 0.35,       # Easy tasks: low threshold, use cheap models
+    "document_drafting": 0.45,   # Medium-easy tasks
+    "tool_heavy": 0.45,         # Tool orchestration, not deep reasoning
+    "retrieval_heavy": 0.45,    # Retrieval-heavy, moderate reasoning
+    "coding": 0.55,             # Coding needs decent models
+    "research": 0.55,           # Research needs good models
+    "unknown_ambiguous": 0.60,  # Unknown = be careful
+    "long_horizon": 0.60,       # Long horizon = be careful
+    "legal_regulated": 0.75,    # Legal = always verify, escalate aggressively
+}
+
+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?",
+        "What is 2+2?","Name the planets in the solar system."],
+    "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.",
+        "Write a fibonacci function with memoization."],
+    "research":["Research latest transformer advances.",
+        "Find sources comparing LoRA and full FT.",
+        "Investigate data center climate impact.",
+        "Survey privacy-preserving ML techniques.",
+        "Compare reinforcement learning algorithms for robotics.",
+        "Analyze recent papers on mixture of experts."],
+    "document_drafting":["Draft project proposal for ML pipeline.",
+        "Write email to team about deployment.","Create technical report on performance.",
+        "Write a project brief for the migration.","Draft meeting agenda."],
+    "legal_regulated":["Review this contract for liability clauses.",
+        "Check GDPR compliance for data pipeline.","Draft privacy policy section.",
+        "Verify regulatory compliance for medical device software.",
+        "Analyze indemnification clause in vendor agreement."],
+    "tool_heavy":["Search open issues and create summary.",
+        "Fetch API docs and generate client code.","Query Q3 sales and produce chart.",
+        "Aggregate logs from 3 services."],
+    "retrieval_heavy":["Answer based on 50-page document.",
+        "Find all payment processing mentions.","Retrieve relevant cases for legal query.",
+        "Summarize findings from quarterly report."],
+    "long_horizon":["Plan 3-month roadmap.","Orchestrate multi-region deployment.",
+        "Redesign data architecture end-to-end.","Migrate monolith to microservices."],
+    "unknown_ambiguous":["Help me with this thing.",
+        "I need something about the server.","Can you look into that issue?",
+        "There's a problem with the data."],
+}
+
+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, oversample_easy=False):
+    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
+    
+    # When oversampling: bias actual_tier toward optimal to create more success examples
+    if oversample_easy and opt <= 2:
+        actual = rng.choices([1,2,3,4,5], weights=[2,5,2,1,0.5])[0]
+    elif oversample_easy and opt <= 3:
+        actual = rng.choices([1,2,3,4,5], weights=[0.5,1,5,2,0.5])[0]
+    else:
+        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 v5: CALIBRATED + PER-TASK THRESHOLDS")
+print("="*80)
+
+# ─── Generate Training Data with Oversampling ────────────────────────
+print("\n[1] Generating 60K training traces (with easy-task oversampling)...")
+rng = random.Random(42)
+
+# Base 50K traces
+traces = [gen_trace(i, rng, oversample_easy=False) for i in range(40000)]
+# Add 20K oversampled easy-task traces
+traces += [gen_trace(i+40000, rng, oversample_easy=True) for i in range(20000)]
+
+print(f"  Total: {len(traces)} traces")
+
+# Check success rate per tier
+for tier in range(1, 6):
+    succ = sum(1 for t in traces if t["tier_out"].get(tier, False))
+    print(f"  Tier {tier}: success rate = {succ/len(traces):.3f}")
+
+# ─── Build Feature Matrix ────────────────────────────────────────────
+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, brier_score_loss
+import xgboost as xgb
+from sklearn.calibration import CalibratedClassifierCV
+
+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"\n  Train: {len(X_train)}, Test: {len(X_test)}")
+
+# ─── Train + Calibrate Per-Tier Classifiers ──────────────────────────
+print("\n[2] Training + calibrating per-tier P(success) classifiers...")
+tier_clfs = {}
+tier_calibrators = {}
+
+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)
+    
+    # Train XGBoost
+    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)
+    
+    # Platt scaling calibration
+    from sklearn.linear_model import LogisticRegression
+    from sklearn.isotonic import IsotonicRegression
+    
+    # Get raw probabilities on test set for calibration
+    y_prob_raw = clf.predict_proba(X_test)[:, 1]
+    
+    # Use isotonic regression for calibration (works better than Platt for small datasets)
+    iso_reg = IsotonicRegression(out_of_bounds="clip")
+    iso_reg.fit(y_prob_raw, y_te)
+    
+    # Evaluate calibration
+    y_prob_cal = iso_reg.transform(y_prob_raw)
+    brier_raw = brier_score_loss(y_te, y_prob_raw)
+    brier_cal = brier_score_loss(y_te, y_prob_cal)
+    
+    acc = accuracy_score(y_te, clf.predict(X_test))
+    f1 = f1_score(y_te, clf.predict(X_test), zero_division=0)
+    
+    tier_clfs[tier] = clf
+    tier_calibrators[tier] = iso_reg
+    print(f"  Tier {tier}: acc={acc:.3f}, f1={f1:.3f}, brier_raw={brier_raw:.3f}, brier_cal={brier_cal:.3f}")
+
+# ─── Calibrated Router ────────────────────────────────────────────────
+print("\n[3] Building calibrated per-task-threshold router...")
+
+def route_calibrated(request, task_type, difficulty):
+    """Calibrated router with per-task thresholds."""
+    base_tier = min(difficulty + 1, 5)
+    floor = TASK_FLOOR.get(task_type, 2)
+    base_tier = max(base_tier, floor)
+    
+    feats = extract_features(request, task_type, difficulty)
+    x = f2v(feats).reshape(1, -1)
+    
+    # Get CALIBRATED P(success) at base_tier
+    p_raw = tier_clfs[base_tier].predict_proba(x)[0, 1]
+    p_success = float(tier_calibrators[base_tier].transform([p_raw])[0])
+    
+    # Per-task threshold
+    threshold = TASK_THRESHOLDS.get(task_type, 0.55)
+    
+    # Escalate if calibrated probability too low
+    while p_success < threshold and base_tier < 5:
+        base_tier += 1
+        p_raw = tier_clfs[base_tier].predict_proba(x)[0, 1]
+        p_success = float(tier_calibrators[base_tier].transform([p_raw])[0])
+    
+    return base_tier
+
+# ─── Evaluate ─────────────────────────────────────────────────────────
+print("\n[4] Evaluating all routers on 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)} eval 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)}
+
+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"])
+# results["v4_prod_t0.55"] = eval_router("v4_prod_t0.55", 
+    # lambda t: route_v4(t, 0.55))
+results["v5_calibrated"] = eval_router("v5_calibrated",
+    lambda t: route_calibrated(t["req"], t["tt"], t["diff"]))
+
+# v4 router for comparison
+def route_v4(t, threshold):
+    base = min(t["diff"]+1, 5)
+    floor = TASK_FLOOR.get(t["tt"], 2)
+    base = max(base, floor)
+    feats = extract_features(t["req"], t["tt"], t["diff"])
+    x = f2v(feats).reshape(1, -1)
+    ps = tier_clfs[base].predict_proba(x)[0, 1]
+    while ps < threshold and base < 5:
+        base += 1
+        ps = tier_clfs[base].predict_proba(x)[0, 1]
+    return base
+
+# Print
+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}")
+
+# Per-task breakdown
+print(f"\n\n[5] Per-task-type breakdown (calibrated v5 vs frontier vs heuristic)...")
+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)),
+                       ("calibrated", lambda t:route_calibrated(t["req"],t["tt"],t["diff"])),
+                       ("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 Frontier ──────────────────────────────────────────────────
+print(f"\n\n[6] Pareto frontier analysis...")
+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}% unsafe={r['unsafe_rate']:.3f}")
+
+# ─── Save Final Production Model ──────────────────────────────────────
+print("\n\n[7] Saving final production model bundle...")
+os.makedirs("/app/router_models", exist_ok=True)
+
+import pickle
+
+bundle = {
+    "tier_clfs": {str(k): v for k, v in tier_clfs.items()},
+    "tier_calibrators": {str(k): v for k, v in tier_calibrators.items()},
+    "feat_keys": FEAT_KEYS,
+    "tier_config": {
+        "tier_cost": TIER_COST,
+        "tier_str": TIER_STR,
+        "task_floor": TASK_FLOOR,
+        "task_thresholds": TASK_THRESHOLDS,
+    },
+    "version": "5.0",
+    "description": "ACO Production Router v5: calibrated + per-task thresholds + oversampled",
+}
+
+with open("/app/router_models/router_bundle_v5.pkl", "wb") as f:
+    pickle.dump(bundle, f)
+print(f"  Saved router_bundle_v5.pkl ({os.path.getsize('/app/router_models/router_bundle_v5.pkl')/1024:.0f} KB)")
+
+# Also save individual files
+for tier in range(1, 6):
+    tier_clfs[tier].save_model(f"/app/router_models/v5_tier_{tier}_success.json")
+with open("/app/router_models/v5_feat_keys.json","w") as f:
+    json.dump(FEAT_KEYS, f)
+with open("/app/router_models/v5_tier_config.json","w") as f:
+    json.dump(bundle["tier_config"], f, indent=2)
+with open("/app/router_models/v5_calibrators.pkl","wb") as f:
+    pickle.dump(tier_calibrators, f)
+
+# Save eval results
+with open("/app/router_models/v5_eval_results.json","w") as f:
+    json.dump(results, f, indent=2, default=str)
+
+print(f"\n\n{'='*80}")
+print("FINAL v5 COMPARISON")
+print(f"{'='*80}")
+print(f"\n{'Router':<25} {'Success':>10} {'AvgCost':>10} {'CostRed':>10} {'Unsafe':>10} {'F-DONE':>10}")
+print("-"*75)
+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}")
+
+print(f"\nDONE!")