Upload training/train_router_v3.py with huggingface_hub

training/train_router_v3.py

@@ -0,0 +1,262 @@
+#!/usr/bin/env python3
+"""Trained Router v3: Combined best approach.
+Uses per-tier P(success) classifiers + safety floors + 
+cost-aware routing with ASYMMETRIC penalties (underkill penalized 5x harder than overkill).
+"""
+import json, os, sys, random, pickle, uuid
+import numpy as np
+from datetime import datetime
+from collections import defaultdict
+from typing import Dict, List, Tuple, Any, Optional
+
+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.",
+        "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?"],
+}
+
+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 v3")
+print("Asymmetric cost scoring: underkill 5x penalty")
+print("="*80)
+
+# ─── Generate ────────────────────────────────────────────────────────
+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
+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)}")
+
+# ─── Train Per-Tier Classifiers ─────────────────────────────────────
+print("\n[2] Training per-tier P(success) classifiers...")
+tier_clfs = {}
+for tier in range(1,6):
+    y_tr = per_tier_labels[tier][idx_train]
+    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(per_tier_labels[tier][idx_test], y_pred)
+    f1 = f1_score(per_tier_labels[tier][idx_test], y_pred, zero_division=0)
+    tier_clfs[tier] = clf
+    print(f"  Tier {tier}: acc={acc:.3f}, f1={f1:.3f}")
+
+# ─── Asymmetric Cost Router ──────────────────────────────────────────
+print("\n[3] Building asymmetric cost router...")
+
+def route_asymmetric(x, task_type, tier_clfs, underkill_penalty=5.0, overkill_penalty=1.0):
+    """Score each tier with asymmetric penalties.
+    
+    score(tier) = P(failure@tier) * underkill_penalty * cost_of_failure
+                + cost(tier) * overkill_penalty
+    
+    Underkill (routing too low) is penalized 5x more than overkill.
+    """
+    if x.ndim == 1:
+        x = x.reshape(1, -1)
+    floor = TASK_FLOOR.get(task_type, 2)
+
+    best_tier = floor
+    best_score = float("inf")
+
+    for tier in range(floor, 6):
+        p_fail = 1.0 - tier_clfs[tier].predict_proba(x)[0, 1]
+        cost_norm = TIER_COST[tier] / TIER_COST[5]  # [0.03, 1.0]
+        
+        # Expected cost of failure (cheap model on hard task)
+        failure_cost = p_fail * underkill_penalty
+        
+        # Cost of using this tier (overkill penalty)
+        tier_cost = cost_norm * overkill_penalty
+        
+        score = failure_cost + tier_cost
+        
+        if score < best_score:
+            best_score = score
+            best_tier = tier
+
+    return best_tier
+
+# ─── Evaluate ─────────────────────────────────────────────────────────
+print("\n[4] Evaluating all routers...")
+n_test = len(idx_test)
+results = {}
+
+def eval_router(name, route_fn):
+    succ = 0; cost = 0.0; unsafe = 0; fd = 0
+    td = defaultdict(int)
+    for i in idx_test:
+        t = traces[i]
+        x = f2v(t["feats"]).reshape(1,-1)
+        pred = route_fn(x, t)
+        td[pred] += 1
+        if t["tier_out"].get(pred, False): succ += 1
+        elif pred < t["opt"]: unsafe += 1
+        elif pred >= t["opt"]: fd += 1
+        cost += TIER_COST[pred]
+    results[name] = {"success":succ/n_test, "avg_cost":cost/n_test,
+                     "unsafe_rate":unsafe/n_test, "false_done":fd/n_test,
+                     "tier_dist":dict(td)}
+
+eval_router("always_frontier", lambda x,t: 4)
+eval_router("always_cheap", lambda x,t: 1)
+eval_router("heuristic_diff+1", lambda x,t: min(t["diff"]+1,5))
+eval_router("heuristic_floor", lambda x,t: TASK_FLOOR.get(t["tt"],2))
+
+for ukp in [3.0, 5.0, 8.0, 10.0, 15.0]:
+    eval_router(f"asymmetric_uk{ukp:.0f}", lambda x,t,uk=ukp: route_asymmetric(x, t["tt"], tier_clfs, underkill_penalty=uk))
+
+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)
+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}")
+
+# Pareto 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}% unsafe={r['unsafe_rate']:.3f}")
+
+# ─── Save Best Model ─────────────────────────────────────────────────
+print("\n[5] 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")
+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)
+with open("/app/router_models/eval_results_v3.json","w") as f:
+    json.dump(results, f, indent=2, default=str)
+print(f"  Saved all models to /app/router_models/")
+
+# Find best config
+best_name = ""
+best_score = -float("inf")
+for name, r in results.items():
+    if name in ("oracle","always_cheap"): continue
+    # Composite: success*20 - cost*50 - unsafe*100
+    score = r["success"]*20 - r["avg_cost"]*50 - r["unsafe_rate"]*100
+    if score > best_score:
+        best_score = score
+        best_name = name
+print(f"\n  BEST CONFIG: {best_name} (composite score: {best_score:.2f})")
+print(f"\nDONE!")