Upload training/train_router_full.py with huggingface_hub

training/train_router_full.py

@@ -0,0 +1,433 @@
+#!/usr/bin/env python3
+"""Train a learned model router for Agent Cost Optimizer."""
+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_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 TRAINING")
+print("="*80)
+
+# ─── Generate Training Data ────────────────────────────────────────
+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")
+
+opt_dist = defaultdict(int)
+for t in traces:
+    opt_dist[t["opt"]] += 1
+for k in sorted(opt_dist):
+    print(f"  opt_tier={k}: {opt_dist[k]} ({opt_dist[k]/len(traces)*100:.1f}%)")
+
+# ─── Build Feature Matrix ──────────────────────────────────────────
+print("\n[2] Building feature matrix...")
+
+def feats_to_vec(feats):
+    """Convert feature dict to fixed-order numpy array."""
+    keys = sorted(feats.keys())
+    return np.array([float(feats[k]) for k in keys], dtype=np.float32)
+
+# Get feature key order from first trace
+FEAT_KEYS = sorted(traces[0]["feats"].keys())
+NUM_FEATURES = len(FEAT_KEYS)
+print(f"  Feature count: {NUM_FEATURES}")
+print(f"  Features: {FEAT_KEYS}")
+
+def feats_to_vec_safe(feats):
+    return np.array([float(feats.get(k, 0.0)) for k in FEAT_KEYS], dtype=np.float32)
+
+X_all = np.array([feats_to_vec_safe(t["feats"]) for t in traces])
+y_opt = np.array([t["opt"] for t in traces])
+y_actual = np.array([t["actual"] for t in traces])
+y_outcome = np.array([1 if t["outcome"]=="success" else 0 for t in traces])
+
+# ─── Per-Tier Success Classifiers ───────────────────────────────────
+print("\n[3] Training per-tier P(success|query) classifiers...")
+from sklearn.model_selection import train_test_split
+from sklearn.metrics import classification_report, accuracy_score, f1_score
+
+# For each tier, create binary label: did this tier succeed?
+per_tier_labels = {}
+for tier in range(1, 6):
+    labels = []
+    for t in traces:
+        labels.append(1 if t["tier_out"].get(tier, False) else 0)
+    per_tier_labels[tier] = np.array(labels)
+    succ_rate = per_tier_labels[tier].mean()
+    print(f"  Tier {tier}: success rate = {succ_rate:.3f}")
+
+# Split train/test
+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)}")
+
+# ─── XGBoost Per-Tier Classifiers ───────────────────────────────────
+print("\n[4] Training XGBoost per-tier classifiers...")
+import xgboost as xgb
+
+tier_classifiers = {}
+for tier in range(1, 6):
+    y_train_tier = per_tier_labels[tier][idx_train]
+    y_test_tier = per_tier_labels[tier][idx_test]
+
+    clf = xgb.XGBClassifier(
+        n_estimators=100,
+        max_depth=4,
+        learning_rate=0.1,
+        subsample=0.8,
+        colsample_bytree=0.8,
+        objective="binary:logistic",
+        eval_metric="logloss",
+        random_state=42,
+        verbosity=0,
+    )
+    clf.fit(X_train, y_train_tier)
+
+    y_pred = clf.predict(X_test)
+    y_prob = clf.predict_proba(X_test)[:, 1]
+
+    acc = accuracy_score(y_test_tier, y_pred)
+    f1 = f1_score(y_test_tier, y_pred, zero_division=0)
+
+    tier_classifiers[tier] = clf
+    print(f"  Tier {tier}: accuracy={acc:.3f}, f1={f1:.3f}")
+
+# ─── CARROT-Style Router Decision ────────────────────────────────────
+print("\n[5] Building CARROT-style router...")
+
+def route_carrot(features_vec, tier_clfs, mu=0.7):
+    """Route to argmin_tier [mu*(1-P_success_tier) + (1-mu)*cost_tier].
+    
+    mu controls quality-vs-cost tradeoff:
+      mu=1.0: maximize quality only (always frontier)
+      mu=0.0: minimize cost only (always cheapest)
+      mu=0.7: 70% quality, 30% cost (our default)
+    """
+    if features_vec.ndim == 1:
+        features_vec = features_vec.reshape(1, -1)
+
+    best_tier = 3
+    best_score = float("inf")
+
+    for tier in range(1, 6):
+        p_success = tier_clfs[tier].predict_proba(features_vec)[0, 1]
+        cost_norm = TIER_COST[tier] / TIER_COST[5]  # normalize to [0,1]
+        score = mu * (1.0 - p_success) + (1.0 - mu) * cost_norm
+        if score < best_score:
+            best_score = score
+            best_tier = tier
+
+    return best_tier
+
+# Evaluate on test set
+print("\n[6] Evaluating CARROT router on test set...")
+
+mu_values = [0.5, 0.6, 0.7, 0.8, 0.9, 1.0]
+
+for mu in mu_values:
+    correct = 0
+    total_cost = 0.0
+    unsafe_misses = 0
+    for i in idx_test:
+        t = traces[i]
+        x = feats_to_vec_safe(t["feats"]).reshape(1, -1)
+        pred_tier = route_carrot(x, tier_classifiers, mu=mu)
+        opt_tier = t["opt"]
+
+        # Check if predicted tier would succeed
+        would_succeed = t["tier_out"].get(pred_tier, False)
+        if would_succeed:
+            correct += 1
+
+        # Cost of predicted tier
+        total_cost += TIER_COST[pred_tier]
+
+        # Unsafe miss: predicted cheap tier for hard task
+        if pred_tier < opt_tier and not would_succeed:
+            unsafe_misses += 1
+
+    n_test = len(idx_test)
+    success_rate = correct / n_test
+    avg_cost = total_cost / n_test
+    unsafe_rate = unsafe_misses / n_test
+
+    # Compare to heuristic (task-type based)
+    heuristic_correct = 0
+    heuristic_cost = 0.0
+    for i in idx_test:
+        t = traces[i]
+        # Heuristic: route by task type (from classifier.py)
+        tt = t["tt"]
+        diff = t["diff"]
+        h_tier = min(diff + 1, 5)  # simple: difficulty + 1
+        if t["tier_out"].get(h_tier, False):
+            heuristic_correct += 1
+        heuristic_cost += TIER_COST[h_tier]
+
+    h_success = heuristic_correct / n_test
+    h_avg_cost = heuristic_cost / n_test
+
+    # Frontier baseline
+    frontier_correct = sum(1 for i in idx_test if traces[i]["tier_out"].get(4, False) or traces[i]["tier_out"].get(5, False))
+    frontier_rate = frontier_correct / n_test
+    frontier_avg_cost = TIER_COST[4]  # always tier 4
+
+    print(f"\n  mu={mu:.1f}:")
+    print(f"    CARROT:    success={success_rate:.3f}, avg_cost={avg_cost:.4f}, unsafe_miss={unsafe_rate:.3f}")
+    print(f"    Heuristic: success={h_success:.3f}, avg_cost={h_avg_cost:.4f}")
+    print(f"    Frontier:  success={frontier_rate:.3f}, avg_cost={frontier_avg_cost:.4f}")
+    print(f"    Cost reduction vs frontier: {(1-avg_cost/frontier_avg_cost)*100:.1f}%")
+    print(f"    Cost reduction vs heuristic: {(1-avg_cost/h_avg_cost)*100:.1f}%")
+
+# ─── XGBoost Direct Optimal-Tier Classifier ─────────────────────────
+print("\n\n[7] Training XGBoost direct optimal-tier classifier...")
+
+y_train_opt = y_opt[idx_train] - 1  # XGB needs 0-indexed
+y_test_opt = y_opt[idx_test] - 1
+
+direct_clf = xgb.XGBClassifier(
+    n_estimators=200,
+    max_depth=5,
+    learning_rate=0.1,
+    subsample=0.8,
+    colsample_bytree=0.8,
+    objective="multi:softmax",
+    num_class=6,
+    eval_metric="mlogloss",
+    random_state=42,
+    verbosity=0,
+)
+direct_clf.fit(X_train, y_train_opt)
+
+y_pred_opt = direct_clf.predict(X_test)
+acc_direct = accuracy_score(y_test_opt, y_pred_opt)
+print(f"  Direct classifier accuracy: {acc_direct:.3f}")
+
+# Detailed classification report
+from sklearn.metrics import confusion_matrix
+cm = confusion_matrix(y_test_opt, y_pred_opt, labels=[1,2,3,4,5])
+print(f"\n  Confusion Matrix (rows=true, cols=predicted):")
+print(f"  {'':>10} {'T1':>6} {'T2':>6} {'T3':>6} {'T4':>6} {'T5':>6}")
+for i, tier in enumerate([1,2,3,4,5]):
+    row = f"  True T{tier:>1}:"
+    for j in range(5):
+        row += f" {cm[i][j]:>6}"
+    print(row)
+
+# Evaluate direct classifier on test set
+print("\n[8] Evaluating direct optimal-tier classifier...")
+direct_correct = 0
+direct_cost = 0.0
+direct_unsafe = 0
+direct_underkill = 0
+direct_overkill = 0
+
+for i, test_idx in enumerate(idx_test):
+    t = traces[test_idx]
+    x = feats_to_vec_safe(t["feats"]).reshape(1, -1)
+    pred_tier = int(direct_clf.predict(x)[0]) + 1  # back to 1-indexed
+    opt_tier = t["opt"]
+
+    would_succeed = t["tier_out"].get(pred_tier, False)
+    if would_succeed:
+        direct_correct += 1
+    direct_cost += TIER_COST[pred_tier]
+
+    if pred_tier < opt_tier:
+        direct_underkill += 1
+        if not would_succeed:
+            direct_unsafe += 1
+    elif pred_tier > opt_tier:
+        direct_overkill += 1
+
+n_test = len(idx_test)
+print(f"  Success rate: {direct_correct/n_test:.3f}")
+print(f"  Avg cost: {direct_cost/n_test:.4f}")
+print(f"  Underkill (pred < optimal): {direct_underkill/n_test:.3f}")
+print(f"  Overkill (pred > optimal): {direct_overkill/n_test:.3f}")
+print(f"  Unsafe misses: {direct_unsafe/n_test:.3f}")
+
+# ─── Comparison: All Routers ────────────────────────────────────────
+print("\n\n" + "="*80)
+print("FINAL COMPARISON: ALL ROUTERS ON TEST SET")
+print("="*80)
+
+# 1. Always frontier
+f_succ = sum(1 for i in idx_test if traces[i]["tier_out"].get(4,False) or traces[i]["tier_out"].get(5,False))
+f_cost = TIER_COST[4] * n_test
+
+# 2. Always cheapest
+c_succ = sum(1 for i in idx_test if traces[i]["tier_out"].get(1,False) or traces[i]["tier_out"].get(2,False))
+c_cost = TIER_COST[1] * n_test
+
+# 3. Heuristic (difficulty + 1)
+h_succ = 0; h_cost = 0.0
+for i in idx_test:
+    t = traces[i]
+    h_tier = min(t["diff"] + 1, 5)
+    if t["tier_out"].get(h_tier, False): h_succ += 1
+    h_cost += TIER_COST[h_tier]
+
+# 4. CARROT (best mu)
+best_mu = 0.7
+carrot_succ = 0; carrot_cost = 0.0; carrot_unsafe = 0
+for i in idx_test:
+    t = traces[i]
+    x = feats_to_vec_safe(t["feats"]).reshape(1, -1)
+    pred = route_carrot(x, tier_classifiers, mu=best_mu)
+    if t["tier_out"].get(pred, False): carrot_succ += 1
+    carrot_cost += TIER_COST[pred]
+    if pred < t["opt"] and not t["tier_out"].get(pred, False):
+        carrot_unsafe += 1
+
+# 5. Direct XGB
+d_succ = direct_correct
+d_cost = direct_cost
+d_unsafe = direct_unsafe
+
+# 6. Oracle (always picks optimal)
+o_succ = sum(1 for i in idx_test if traces[i]["tier_out"].get(traces[i]["opt"], False))
+o_cost = sum(TIER_COST[traces[i]["opt"]] for i in idx_test)
+
+print(f"\n{'Router':<20} {'Success':>10} {'AvgCost':>10} {'CostRed':>10} {'Unsafe':>10}")
+print("-"*60)
+for name, succ, cost, unsafe in [
+    ("always_frontier", f_succ, f_cost, 0),
+    ("always_cheap", c_succ, c_cost, 0),
+    ("heuristic", h_succ, h_cost, 0),
+    (f"CARROT(mu={best_mu})", carrot_succ, carrot_cost, carrot_unsafe),
+    ("direct_xgb", d_succ, d_cost, d_unsafe),
+    ("oracle", o_succ, o_cost, 0),
+]:
+    sr = succ/n_test
+    ac = cost/n_test
+    cr = (1 - cost/f_cost)*100
+    um = unsafe/n_test
+    print(f"{name:<20} {sr:>10.3f} {ac:>10.4f} {cr:>9.1f}% {um:>10.3f}")
+
+# ─── Save Models ────────────────────────────────────────────────────
+print("\n\n[9] Saving models...")
+os.makedirs("/app/router_models", exist_ok=True)
+
+# Save per-tier classifiers
+for tier, clf in tier_classifiers.items():
+    clf.save_model(f"/app/router_models/tier_{tier}_success.json")
+    print(f"  Saved tier_{tier}_success.json")
+
+# Save direct classifier
+direct_clf.save_model("/app/router_models/direct_optimal_tier.json")
+print(f"  Saved direct_optimal_tier.json")
+
+# Save feature keys
+with open("/app/router_models/feat_keys.json", "w") as f:
+    json.dump(FEAT_KEYS, f)
+print(f"  Saved feat_keys.json ({len(FEAT_KEYS)} features)")
+
+# Save tier config
+with open("/app/router_models/tier_config.json", "w") as f:
+    json.dump({"tier_cost": TIER_COST, "tier_str": TIER_STR}, f)
+print(f"  Saved tier_config.json")
+
+print("\n\nDONE! Models saved to /app/router_models/")
+print("Next step: integrate trained router into ACO ModelCascadeRouter._route_learned()")