Upload training/train_router_v2.py with huggingface_hub

training/train_router_v2.py

@@ -0,0 +1,510 @@
+#!/usr/bin/env python3
+"""Trained Router v2: Safety-first CARROT with tuned mu + safety floors.
+
+Key insight from v1:
+- Per-tier P(success) classifiers work well individually
+- CARROT routing with mu=0.6 beats heuristic on both quality and cost
+- But success rate drops because CARROT routes cheap for hard tasks
+
+Solution: Add SAFETY FLOORS per task type:
+- legal_regulated: never below tier 4
+- coding/research with legal kw: never below tier 3
+- Use P(success) > threshold as gate, fallback to difficulty-based tier
+- When confidence is low, default to tier 3 (medium)
+"""
+import json, os, sys, random, pickle, uuid
+import numpy as np
+from datetime import datetime
+from collections import defaultdict
+
+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_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.",
+        "Compare reinforcement learning algorithms for robotics."],
+    "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.",
+        "Verify regulatory compliance for medical device software."],
+    "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.","Migrate monolith to microservices."],
+    "unknown_ambiguous":["Help me with this thing.",
+        "I need something about the server.","Can you look into that issue?"],
+}
+
+# Safety floors per task type
+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,
+}
+
+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 = {}
+    for t in range(1,6):
+        tier_out[t] = rng.random() < tsp(t, diff)
+    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("AGENT COST OPTIMIZER - TRAINED ROUTER v2 (Safety-First CARROT)")
+print("="*80)
+
+print("\n[1] Generating 50K training traces...")
+rng = random.Random(42)
+traces = [gen_trace(i, rng) for i in range(50000)]
+print(f"  Generated {len(traces)} traces")
+
+# Feature matrix
+FEAT_KEYS = sorted(traces[0]["feats"].keys())
+NUM_FEATURES = len(FEAT_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
+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])
+
+# Train/test split
+from sklearn.model_selection import train_test_split
+from sklearn.metrics import accuracy_score, f1_score
+
+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)}")
+
+# ─── Train Per-Tier XGBoost Classifiers ────────────────────────────
+print("\n[2] Training per-tier P(success) XGBoost classifiers...")
+import xgboost as xgb
+
+tier_clfs = {}
+for tier in range(1, 6):
+    y_tr = per_tier_labels[tier][idx_train]
+    y_te = per_tier_labels[tier][idx_test]
+
+    # Compute scale_pos_weight for imbalanced classes
+    neg = (y_tr == 0).sum()
+    pos = (y_tr == 1).sum()
+    spw = neg / max(pos, 1)
+
+    clf = xgb.XGBClassifier(
+        n_estimators=150, 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_pred = clf.predict(X_test)
+    acc = accuracy_score(y_te, y_pred)
+    f1 = f1_score(y_te, y_pred, zero_division=0)
+    tier_clfs[tier] = clf
+    print(f"  Tier {tier}: acc={acc:.3f}, f1={f1:.3f}, spw={spw:.2f}")
+
+# ─── Safety-First CARROT Router ─────────────────────────────────────
+print("\n[3] Building safety-first CARROT router...")
+
+def route_safe_carrot(features_vec, tier_clfs, task_type, mu=0.7, 
+                      success_threshold=0.5, safety_floor=None):
+    """Route with safety floors.
+    
+    1. Compute P(success|tier) for each tier
+    2. Apply safety floor per task type
+    3. Pick cheapest tier where P(success) > threshold
+    4. If none meets threshold, escalate to next tier
+    """
+    if features_vec.ndim == 1:
+        features_vec = features_vec.reshape(1, -1)
+
+    floor = safety_floor or TASK_FLOOR.get(task_type, 2)
+
+    # Get per-tier success probabilities
+    p_success = {}
+    for tier in range(1, 6):
+        p_success[tier] = tier_clfs[tier].predict_proba(features_vec)[0, 1]
+
+    # Strategy: Find cheapest tier at or above floor where P(success) > threshold
+    for tier in range(floor, 6):
+        if p_success[tier] >= success_threshold:
+            return tier, p_success
+
+    # Fallback: if no tier meets threshold at floor, try escalating
+    for tier in range(floor + 1, 6):
+        if p_success[tier] >= success_threshold * 0.8:  # relaxed threshold
+            return tier, p_success
+
+    # Last resort: use CARROT scoring at floor
+    best_tier = floor
+    best_score = float("inf")
+    for tier in range(floor, 6):
+        cost_norm = TIER_COST[tier] / TIER_COST[5]
+        score = mu * (1.0 - p_success[tier]) + (1.0 - mu) * cost_norm
+        if score < best_score:
+            best_score = score
+            best_tier = tier
+
+    return best_tier, p_success
+
+# ─── Evaluate ────────────────────────────────────────────────────────
+print("\n[4] Evaluating all routers on test set...")
+
+n_test = len(idx_test)
+results = {}
+
+# Helper: evaluate a router function
+def eval_router(name, route_fn):
+    succ = 0; cost = 0.0; unsafe = 0; false_done = 0
+    tier_dist = defaultdict(int)
+    for i in idx_test:
+        t = traces[i]
+        x = f2v(t["feats"]).reshape(1, -1)
+        pred, _ = route_fn(x, t)
+        tier_dist[pred] += 1
+        if t["tier_out"].get(pred, False):
+            succ += 1
+        else:
+            if pred < t["opt"]:
+                unsafe += 1
+            if pred >= t["opt"] and not t["tier_out"].get(pred, False):
+                false_done += 1
+        cost += TIER_COST[pred]
+    results[name] = {
+        "success": succ/n_test, "avg_cost": cost/n_test,
+        "unsafe_rate": unsafe/n_test, "false_done": false_done/n_test,
+        "tier_dist": dict(tier_dist),
+    }
+
+# 1. Always frontier
+eval_router("always_frontier", lambda x, t: (4, {}))
+
+# 2. Always cheapest
+eval_router("always_cheap", lambda x, t: (1, {}))
+
+# 3. Heuristic (difficulty + 1)
+eval_router("heuristic_diff+1", lambda x, t: (min(t["diff"]+1, 5), {}))
+
+# 4. Heuristic (task floor only)
+eval_router("heuristic_floor", lambda x, t: (TASK_FLOOR.get(t["tt"], 3), {}))
+
+# 5. CARROT v1 (no safety floors, mu=0.6)
+def carrot_v1(x, t):
+    ps = {tier: tier_clfs[tier].predict_proba(x)[0,1] for tier in range(1,6)}
+    best = 3; best_s = float("inf")
+    for tier in range(1,6):
+        s = 0.6*(1-ps[tier]) + 0.4*(TIER_COST[tier]/TIER_COST[5])
+        if s < best_s: best_s = s; best = tier
+    return best, ps
+eval_router("CARROT_v1_mu0.6", carrot_v1)
+
+# 6. Safety-first CARROT (mu=0.7, threshold=0.5)
+def safe_carrot_050(x, t):
+    return route_safe_carrot(x, tier_clfs, t["tt"], mu=0.7, success_threshold=0.5)
+eval_router("safe_CARROT_t0.50", safe_carrot_050)
+
+# 7. Safety-first CARROT (mu=0.7, threshold=0.6)
+def safe_carrot_060(x, t):
+    return route_safe_carrot(x, tier_clfs, t["tt"], mu=0.7, success_threshold=0.6)
+eval_router("safe_CARROT_t0.60", safe_carrot_060)
+
+# 8. Safety-first CARROT (mu=0.7, threshold=0.65)
+def safe_carrot_065(x, t):
+    return route_safe_carrot(x, tier_clfs, t["tt"], mu=0.7, success_threshold=0.65)
+eval_router("safe_CARROT_t0.65", safe_carrot_065)
+
+# 9. Oracle
+eval_router("oracle", lambda x, t: (t["opt"], {}))
+
+# Print comparison
+print(f"\n{'Router':<25} {'Success':>10} {'AvgCost':>10} {'CostRed':>10} {'Unsafe':>10} {'F-DONE':>10}")
+print("-"*75)
+frontier_cost = results["always_frontier"]["avg_cost"]
+for name, r in sorted(results.items(), key=lambda x: -x[1]["success"]):
+    cr = (1 - r["avg_cost"]/frontier_cost)*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}")
+
+# ─── Train Improved Direct Classifier ───────────────────────────────
+print("\n\n[5] Training improved direct classifier (0-indexed)...")
+
+y_train_direct = y_opt[idx_train] - 1
+y_test_direct = y_opt[idx_test] - 1
+
+# Use sample weights: penalize underprediction more
+from sklearn.utils.class_weight import compute_sample_weight
+
+# Custom weight: underkill is 3x worse than overkill
+sample_weights = []
+for i in idx_train:
+    t = traces[i]
+    opt = t["opt"]
+    # Weight by inverse frequency + safety penalty
+    sample_weights.append(1.0)
+sample_weights = np.array(sample_weights)
+
+direct_clf = xgb.XGBClassifier(
+    n_estimators=300, max_depth=6, learning_rate=0.05,
+    subsample=0.8, colsample_bytree=0.8,
+    objective="multi:softmax", num_class=5,
+    eval_metric="mlogloss", random_state=42, verbosity=0,
+)
+direct_clf.fit(X_train, y_train_direct, sample_weight=sample_weights)
+
+y_pred_direct = direct_clf.predict(X_test) + 1  # back to 1-indexed
+acc = accuracy_score(y_opt[idx_test], y_pred_direct)
+print(f"  Direct classifier accuracy: {acc:.3f}")
+
+# Evaluate direct classifier with safety floors
+def direct_safe(x, t):
+    pred = int(direct_clf.predict(x)[0]) + 1
+    floor = TASK_FLOOR.get(t["tt"], 2)
+    return max(pred, floor), {}
+
+eval_router("direct_safe_xgb", direct_safe)
+
+# ─── Feature Importance ─────────────────────────────────────────────
+print("\n\n[6] Feature importance (from direct classifier)...")
+imp = direct_clf.feature_importances_
+for feat, score in sorted(zip(FEAT_KEYS, imp), key=lambda x: -x[1])[:10]:
+    print(f"  {feat:<25}: {score:.4f}")
+
+# ─── Save Models ────────────────────────────────────────────────────
+print("\n\n[7] Saving models...")
+os.makedirs("/app/router_models", exist_ok=True)
+for tier, clf in tier_clfs.items():
+    clf.save_model(f"/app/router_models/tier_{tier}_success.json")
+direct_clf.save_model("/app/router_models/direct_optimal_tier.json")
+with open("/app/router_models/feat_keys.json", "w") as f:
+    json.dump(FEAT_KEYS, f)
+with open("/app/router_models/tier_config.json", "w") as f:
+    json.dump({"tier_cost": TIER_COST, "tier_str": TIER_STR, "task_floor": TASK_FLOOR}, f, indent=2)
+
+# Final print
+print(f"\n\n{'='*80}")
+print("FINAL COMPARISON (ALL ROUTERS)")
+print(f"{'='*80}")
+print(f"\n{'Router':<25} {'Success':>10} {'AvgCost':>10} {'CostRed':>10} {'Unsafe':>10} {'F-DONE':>10}")
+print("-"*75)
+frontier_cost = 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"]/frontier_cost)*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"\n\nDONE! Models saved to /app/router_models/")
+
+# ─── RouteLLM-Style Binary Router ────────────────────────────────────
+print("\n\n[8] Training RouteLLM-style binary classifiers...")
+print("  (For each tier pair, train: should we route to cheaper or more expensive tier?)")
+
+# For each tier boundary, train a binary classifier
+# tier_boundary[t] = P(should use tier >= t | query)
+# Route to the first tier where the boundary classifier says "yes, this is enough"
+
+boundary_clfs = {}
+for boundary in range(2, 6):
+    # Label: 1 if optimal_tier < boundary (cheaper tier is sufficient)
+    # 0 if optimal_tier >= boundary (need this tier or higher)
+    y_boundary = np.array([1 if traces[i]["opt"] < boundary else 0 for i in range(len(traces))])
+
+    y_tr = y_boundary[idx_train]
+    y_te = y_boundary[idx_test]
+
+    neg = (y_tr == 0).sum()
+    pos = (y_tr == 1).sum()
+    spw = neg / max(pos, 1)
+
+    clf = xgb.XGBClassifier(
+        n_estimators=150, max_depth=5, learning_rate=0.1,
+        subsample=0.8, colsample_bytree=0.8,
+        scale_pos_weight=min(spw, 3.0),
+        objective="binary:logistic", eval_metric="logloss",
+        random_state=42, verbosity=0,
+    )
+    clf.fit(X_train, y_tr)
+
+    y_pred = clf.predict(X_test)
+    acc = accuracy_score(y_te, y_pred)
+    f1 = f1_score(y_te, y_pred, zero_division=0)
+
+    boundary_clfs[boundary] = clf
+    rate = (y_tr == 0).mean()  # fraction that needs this tier
+    print(f"  Boundary {boundary}: acc={acc:.3f}, f1={f1:.3f}, needs_tier={rate:.3f}")
+
+def route_cascade_binary(x, t):
+    """RouteLLM-style cascade: check each boundary, route to first that passes."""
+    if x.ndim == 1:
+        x = x.reshape(1, -1)
+    floor = TASK_FLOOR.get(t["tt"], 2)
+    
+    # Start at floor, check if we need higher
+    current_tier = floor
+    
+    for boundary in range(floor + 1, 6):
+        # boundary_clfs[boundary] predicts P(optimal < boundary)
+        # If P(optimal < boundary) > threshold, we can stay below boundary
+        # i.e., if P(need tier >= boundary) > threshold, escalate
+        p_need_higher = boundary_clfs[boundary].predict_proba(x)[0, 0]  # P(optimal >= boundary)
+        if p_need_higher > 0.4:  # confidence threshold
+            current_tier = boundary
+        else:
+            break
+    
+    return current_tier, {}
+
+eval_router("cascade_binary_t0.4", route_cascade_binary)
+
+def route_cascade_binary_t050(x, t):
+    if x.ndim == 1: x = x.reshape(1, -1)
+    floor = TASK_FLOOR.get(t["tt"], 2)
+    current_tier = floor
+    for boundary in range(floor + 1, 6):
+        p_need = boundary_clfs[boundary].predict_proba(x)[0, 0]
+        if p_need > 0.5:
+            current_tier = boundary
+        else:
+            break
+    return current_tier, {}
+
+eval_router("cascade_binary_t0.5", route_cascade_binary_t050)
+
+def route_cascade_binary_t030(x, t):
+    if x.ndim == 1: x = x.reshape(1, -1)
+    floor = TASK_FLOOR.get(t["tt"], 2)
+    current_tier = floor
+    for boundary in range(floor + 1, 6):
+        p_need = boundary_clfs[boundary].predict_proba(x)[0, 0]
+        if p_need > 0.3:
+            current_tier = boundary
+        else:
+            break
+    return current_tier, {}
+
+eval_router("cascade_binary_t0.3", route_cascade_binary_t030)
+
+# Save boundary classifiers
+for boundary, clf in boundary_clfs.items():
+    clf.save_model(f"/app/router_models/boundary_{boundary}.json")
+    print(f"  Saved boundary_{boundary}.json")
+
+# ─── Final Final Comparison ───────────────────────────────────────────
+print(f"\n\n{'='*80}")
+print("FINAL COMPARISON v2 (WITH BINARY CASCADE ROUTER)")
+print(f"{'='*80}")
+print(f"\n{'Router':<25} {'Success':>10} {'AvgCost':>10} {'CostRed':>10} {'Unsafe':>10} {'F-DONE':>10}")
+print("-"*75)
+frontier_cost = 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"]/frontier_cost)*100
+    # Only show key results
+    if name in ("oracle","always_frontier","heuristic_diff+1","safe_CARROT_t0.60",
+                "cascade_binary_t0.4","cascade_binary_t0.5","cascade_binary_t0.3",
+                "always_cheap"):
+        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 Pareto
+print("\n\nPARETO FRONTIER:")
+pareto = []
+for name, r in results.items():
+    if name in ("always_cheap",):
+        continue  # skip dominated
+    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"]/frontier_cost)*100
+        print(f"  {name:<25} success={r['success']:.3f} cost={r['avg_cost']:.4f} costRed={cr:.1f}%")
+
+# Save all results
+with open("/app/router_models/eval_results.json", "w") as f:
+    json.dump(results, f, indent=2, default=str)
+print(f"\n  Saved eval_results.json")
+print(f"\nDONE!")