Upload ablation_study.py

ablation_study.py

@@ -14,18 +14,16 @@ ACO Ablation Study + Cost-Quality Frontier Report.
  10. no telemetry feedback
 
 Each ablation removes one module from the full ACO and re-runs the benchmark.
-Reports which modules actually save money and which are noise.
+Also runs all 9 baseline configs (A-I) for the frontier report.
 """
 import json, sys, random, os
 from dataclasses import dataclass, field, asdict
 from collections import defaultdict
 from typing import List, Dict
 
-# Import the benchmark suite
 sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)) or ".")
 import importlib.util
 
-# Try local import, fall back to Hub download
 try:
     from benchmark_suite import (
         generate_tasks, simulate_task, Config, CONFIGS, MODELS,
@@ -47,12 +45,7 @@ except ImportError:
 
 random.seed(42)
 
-# ═══════════════════════════════════════════════════════════════════
-# Ablation configs: full ACO minus one module
-# ═══════════════════════════════════════════════════════════════════
-
 def make_ablation_configs() -> List[Config]:
-    """Create 10 ablation configs, each removing one module from full ACO."""
     base = dict(
         use_model_routing=True, use_learned_router=True,
         use_context_budget=True, use_cache_layout=True,
@@ -60,7 +53,7 @@ def make_ablation_configs() -> List[Config]:
         use_retry_optimizer=True, use_meta_tools=True,
         use_early_termination=True, use_telemetry=True,
     )
-    ablations = [
+    return [
         Config("abl1", "no model router", use_model_routing=False, use_learned_router=False,
                use_context_budget=True, use_cache_layout=True, use_tool_gate=True,
                use_verifier_budget=True, use_retry_optimizer=True, use_meta_tools=True,
@@ -79,47 +72,27 @@ def make_ablation_configs() -> List[Config]:
                use_early_termination=True, use_telemetry=True),
         Config("abl10", "no telemetry feedback", **{**base, "use_telemetry": False}),
     ]
-    return ablations
-
-
-# ═══════════════════════════════════════════════════════════════════
-# Override select_model for ablation 9 (no specialist models)
-# ═══════════════════════════════════════════════════════════════════
-
-def select_model_abl9(task) -> str:
-    """No specialist models: always use frontier-tier (tier 4) for routed tasks."""
-    if task.risk_level == "high" or task.difficulty > 0.5:
-        return "gpt-5.2"  # tier 4
-    if task.difficulty > 0.2:
-        return "gpt-5.2"
-    return "gpt-5-mini"  # tier 2 minimum, no tier-1 specialist
-
 
 def run_ablation(n_per_domain: int = 20) -> Dict:
-    """Run all 10 ablations + full ACO baseline."""
+    """Run all baselines + 10 ablations."""
     tasks = generate_tasks(n_per_domain)
     full_aco = Config("I", "full ACO", use_model_routing=True, use_learned_router=True,
                       use_context_budget=True, use_cache_layout=True, use_tool_gate=True,
                       use_verifier_budget=True, use_retry_optimizer=True, use_meta_tools=True,
                       use_early_termination=True, use_telemetry=True)
-    all_configs = [full_aco] + make_ablation_configs()
+    all_configs = list(CONFIGS) + make_ablation_configs()
 
     print(f"Generated {len(tasks)} tasks across 5 domains")
-    print(f"Running {len(all_configs)} configs (full ACO + 10 ablations) x {len(tasks)} tasks\n")
+    print(f"Running {len(all_configs)} configs (9 baselines + 10 ablations) x {len(tasks)} tasks\n")
 
     all_results = []
     for config in all_configs:
         print(f"  {config.name}: {config.label}...", end=" ", flush=True)
         for task in tasks:
-            # Special handling for ablation 9
             if config.name == "abl9":
-                # Override model selection to avoid specialist (tier 1) models
                 result = simulate_task(config, task)
-                # Force model to non-specialist
                 if result["tier"] == 1:
-                    result["model"] = "gpt-5-mini"
-                    result["tier"] = 2
-                    # Recalculate cost
+                    result["model"] = "gpt-5-mini"; result["tier"] = 2
                     mi = MODELS["gpt-5-mini"]
                     result["cost"] = round(
                         (result["input_tokens"] / 1_000_000) * mi["cost_in"] +
@@ -133,7 +106,6 @@ def run_ablation(n_per_domain: int = 20) -> Dict:
 
     return {"tasks": [asdict(t) for t in tasks], "results": all_results}
 
-
 def print_ablation_report(metrics: Dict, config_labels: Dict):
     print(f"\n{'='*100}")
     print(f"  ACO ABLATION REPORT - Which Modules Actually Save Money?")
@@ -146,26 +118,21 @@ def print_ablation_report(metrics: Dict, config_labels: Dict):
     print("-" * 100)
 
     verdicts = []
-    for abl_name in ["abl1", "abl2", "abl3", "abl4", "abl5", "abl6", "abl7", "abl8", "abl9", "abl10"]:
+    for abl_name in [f"abl{i}" for i in range(1, 11)]:
+        if abl_name not in metrics["by_config"]:
+            continue
         m = metrics["by_config"][abl_name]
         label = config_labels.get(abl_name, abl_name)
         cost_delta = m["total_cost"] - full["total_cost"]
         cost_pct = (cost_delta / full["total_cost"]) * 100 if full["total_cost"] > 0 else 0
         quality_delta = (m["success_rate"] - full["success_rate"]) * 100
 
-        # Verdict: does removing this module hurt?
-        if quality_delta < -3:
-            verdict = "CRITICAL"
-        elif cost_delta > 0 and quality_delta < -1:
-            verdict = "HURTS QUALITY"
-        elif cost_delta > 0.02:
-            verdict = "SAVES MONEY"
-        elif abs(cost_pct) < 2 and abs(quality_delta) < 1:
-            verdict = "NOISE"
-        elif cost_delta < 0 and quality_delta >= 0:
-            verdict = "COST INCREASE"
-        else:
-            verdict = "MARGINAL"
+        if quality_delta < -3: verdict = "CRITICAL"
+        elif cost_delta > 0 and quality_delta < -1: verdict = "HURTS QUALITY"
+        elif cost_delta > 0.02: verdict = "SAVES MONEY"
+        elif abs(cost_pct) < 2 and abs(quality_delta) < 1: verdict = "NOISE"
+        elif cost_delta < 0 and quality_delta >= 0: verdict = "COST INCREASE"
+        else: verdict = "MARGINAL"
 
         verdicts.append((abl_name, label, verdict, cost_pct, quality_delta))
         print(f"  {label:<38} {m['success_rate']*100:>6.1f}%  ${m['total_cost']:>8.4f}  "
@@ -181,98 +148,83 @@ def print_ablation_report(metrics: Dict, config_labels: Dict):
     hurts = [v for v in verdicts if v[2] == "HURTS QUALITY"]
     marginal = [v for v in verdicts if v[2] == "MARGINAL"]
 
-    print(f"\n  CRITICAL modules (removing causes >3pp quality loss):")
-    for _, label, _, _, qd in critical:
-        print(f"    - {label}")
-    if not critical: print(f"    (none)")
+    print(f"\n  CRITICAL (removing causes >3pp quality loss):")
+    for _, label, _, _, qd in critical: print(f"    - {label} ({qd:+.1f}pp)")
+    if not critical: print("    (none)")
 
-    print(f"\n  SAVES MONEY modules (removing increases cost):")
-    for _, label, _, cp, _ in saves:
-        print(f"    - {label} (+{cp:.1f}% cost without it)")
-    if not saves: print(f"    (none)")
+    print(f"\n  SAVES MONEY (removing increases cost):")
+    for _, label, _, cp, _ in saves: print(f"    - {label} (+{cp:.1f}% cost without it)")
+    if not saves: print("    (none)")
 
-    print(f"\n  NOISE modules (removing has <2% cost and <1pp quality impact):")
-    for _, label, _, _, _ in noise:
-        print(f"    - {label}")
-    if not noise: print(f"    (none)")
+    print(f"\n  NOISE (removing has <2% cost, <1pp quality):")
+    for _, label, _, _, _ in noise: print(f"    - {label}")
+    if not noise: print("    (none)")
 
-    print(f"\n  HURTS QUALITY modules (removing reduces quality but saves cost):")
-    for _, label, _, _, qd in hurts:
-        print(f"    - {label} ({qd:+.1f}pp quality loss)")
-    if not hurts: print(f"    (none)")
+    print(f"\n  HURTS QUALITY (removing reduces quality but saves cost):")
+    for _, label, _, _, qd in hurts: print(f"    - {label} ({qd:+.1f}pp)")
+    if not hurts: print("    (none)")
 
+    print(f"\n  MARGINAL (small effect either way):")
+    for _, label, _, _, _ in marginal: print(f"    - {label}")
+    if not marginal: print("    (none)")
 
 def print_frontier_report(metrics: Dict, config_labels: Dict):
-    """Cost-quality frontier: plot all configs on cost vs quality axis."""
     print(f"\n{'='*100}")
     print(f"  COST-QUALITY FRONTIER")
     print(f"{'='*100}")
 
+    # Use all configs that exist in metrics
+    available = [k for k in ["A","B","C","D","E","F","G","H","I"] if k in metrics["by_config"]]
     configs = []
-    for cn in ["A", "B", "C", "D", "E", "F", "G", "H", "I"]:
+    for cn in available:
         m = metrics["by_config"][cn]
         configs.append((cn, config_labels.get(cn, cn), m["success_rate"], m["total_cost"]))
 
-    # Sort by cost
     configs.sort(key=lambda x: x[3])
 
-    print(f"\n  {'Config':<40} {'Quality':>8} {'Cost':>10} {'Position':>30}")
-    print("-" * 90)
+    print(f"\n  {'Config':<40} {'Quality':>8} {'Cost':>10}  {'Visual':>35}")
+    print("-" * 95)
 
+    max_cost = max(c[3] for c in configs) if configs else 1
     for cn, label, sr, cost in configs:
-        # Visual bar
         bar_len = int(sr * 30)
-        cost_bar = int(cost / max(c[3] for c in configs) * 20)
         bar = "█" * bar_len + "░" * (30 - bar_len)
         print(f"  {cn}. {label:<36} {sr*100:>6.1f}%  ${cost:>8.4f}  {bar}")
 
-    # Find Pareto-optimal configs
-    print(f"\n  Pareto-optimal configs (no other config is both cheaper AND better):")
     pareto = []
     for i, (cn, label, sr, cost) in enumerate(configs):
         dominated = False
         for j, (cn2, label2, sr2, cost2) in enumerate(configs):
             if i != j and sr2 >= sr and cost2 <= cost and (sr2 > sr or cost2 < cost):
-                dominated = True
-                break
-        if not dominated:
-            pareto.append((cn, label, sr, cost))
+                dominated = True; break
+        if not dominated: pareto.append((cn, label, sr, cost))
 
+    print(f"\n  Pareto-optimal (no other config is both cheaper AND better):")
     for cn, label, sr, cost in pareto:
-        print(f"    {cn}. {label}: {sr*100:.1f}% quality at ${cost:.4f}")
+        print(f"    {cn}. {label}: {sr*100:.1f}% at ${cost:.4f}")
 
-    # Iso-quality analysis
-    print(f"\n  Iso-quality analysis (configs within ±2pp of frontier quality):")
-    frontier_q = max(c[2] for c in configs)
+    frontier_q = max(c[2] for c in configs) if configs else 0
+    print(f"\n  Iso-quality configs (within ±2pp of best quality {frontier_q*100:.1f}%):")
     for cn, label, sr, cost in configs:
         if sr >= frontier_q - 0.02:
-            savings = (1 - cost / max(c[3] for c in configs)) * 100
+            savings = (1 - cost / max_cost) * 100
             print(f"    {cn}. {label}: {sr*100:.1f}% at ${cost:.4f} ({savings:+.1f}% vs most expensive)")
 
-
 def main():
     n = int(sys.argv[1]) if len(sys.argv) > 1 else 20
     data = run_ablation(n)
     metrics = compute_metrics(data["results"])
 
     config_labels = {c.name: c.label for c in CONFIGS}
-    # Add ablation labels
-    for c in make_ablation_configs():
-        config_labels[c.name] = c.label
+    for c in make_ablation_configs(): config_labels[c.name] = c.label
 
     print_ablation_report(metrics, config_labels)
     print_frontier_report(metrics, config_labels)
 
-    output = {
-        "n_tasks_per_domain": n,
-        "metrics": metrics,
-        "config_labels": config_labels,
-        "raw_results": data["results"],
-    }
+    output = {"n_tasks_per_domain": n, "metrics": metrics,
+              "config_labels": config_labels, "raw_results": data["results"]}
     with open("/tmp/aco_ablation_results.json", "w") as f:
         json.dump(output, f, indent=2)
     print(f"\nResults saved to /tmp/aco_ablation_results.json")
 
-
-if __name__ == "__main__":
-    main()
+if __name__ == "__main__": main()