Upload verify_v2.py

verify_v2.py

@@ -1,18 +1,20 @@
 """Independent verification of v2 ACO specialist models.
 
-Run after training completes. Recreates test splits with same seed,
-loads models, computes metrics, compares to v1 baseline.
+Recreates test splits with same seed, loads v1 and v2 models,
+computes metrics, compares delta. Robust to dataset format quirks.
 
-Usage:
-    uv run --with transformers,torch,datasets,scikit-learn,huggingface_hub verify_v2.py
-
-Or via hf_jobs:
-    hf_jobs run --script verify_v2.py --deps transformers,torch,datasets,scikit-learn --hardware a10g-large --timeout 2h
+Usage via hf_jobs:
+    hf_jobs run --script verify_v2.py --deps transformers,torch,datasets,scikit-learn,huggingface_hub --hardware a10g-large --timeout 2h
 """
-import torch, numpy as np, json, os
+import torch, numpy as np, json, os, sys
 from datasets import Dataset, load_dataset
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
-from sklearn.metrics import accuracy_score, f1_score, classification_report
+from sklearn.metrics import accuracy_score, f1_score, classification_report, precision_recall_fscore_support
+from torch.utils.data import DataLoader
+
+# ═══════════════════════════════════════════
+# Constants
+# ═══════════════════════════════════════════
 
 V1_MODELS = {
     "tier_router": "narcolepticchicken/aco-specialists-tier-router",
@@ -25,8 +27,12 @@ V2_MODELS = {
     "verifier_gater": "narcolepticchicken/aco-specialists-verifier-gater-v2",
 }
 NUM_LABELS_MAP = {"tier_router": 3, "tool_gater": 2, "verifier_gater": 2}
+TASK_NAMES = ["tier_router", "tool_gater", "verifier_gater"]
+
+# ═══════════════════════════════════════════
+# Dataset loaders (SAME as training script)
+# ═══════════════════════════════════════════
 
-# --- Dataset loaders (SAME as training script) ---
 def load_tool_gater():
     import re
     ds = load_dataset("lockon/ToolACE", split="train")
@@ -103,21 +109,46 @@ def load_verifier_gater():
 
 LOADERS = {"tool_gater": load_tool_gater, "tier_router": load_tier_router, "verifier_gater": load_verifier_gater}
 
+# ═══════════════════════════════════════════
+# Simple tensor dataset wrapper
+# ═══════════════════════════════════════════
+
+class TensorDataset(torch.utils.data.Dataset):
+    def __init__(self, encodings, labels):
+        self.input_ids = encodings["input_ids"]
+        self.attention_mask = encodings["attention_mask"]
+        self.labels = labels
+
+    def __len__(self):
+        return len(self.labels)
+
+    def __getitem__(self, idx):
+        return {
+            "input_ids": torch.tensor(self.input_ids[idx]),
+            "attention_mask": torch.tensor(self.attention_mask[idx]),
+            "label": torch.tensor(self.labels[idx]),
+        }
+
+# ═══════════════════════════════════════════
+# Evaluation
+# ═══════════════════════════════════════════
+
 def evaluate_model(model_name, task_name, num_labels):
-    """Load model, tokenize test set, compute metrics."""
-    print(f"\n{'='*60}\nEVALUATING: {model_name} [{task_name}]\n{'='*60}")
-    
+    print(f"\n{'='*60}")
+    print(f"EVALUATING: {model_name} [{task_name}]")
+    print(f"{'='*60}")
+
     # Load data
     ds = LOADERS[task_name]()
     test_ds = ds["test"]
     print(f"  Test samples: {len(test_ds)}")
-    
+
     # Class distribution
     lc = {}
     for lb in test_ds["labels"]:
         lc[lb] = lc.get(lb, 0) + 1
     print(f"  Class dist: {lc}")
-    
+
     # Load model
     try:
         tokenizer = AutoTokenizer.from_pretrained(model_name)
@@ -131,85 +162,152 @@ def evaluate_model(model_name, task_name, num_labels):
         print(f"  Model loaded on {device}")
     except Exception as e:
         print(f"  FAILED to load model: {e}")
+        import traceback; traceback.print_exc()
         return None
-    
+
     # Extract threshold from config
     threshold = getattr(model.config, "threshold", 0.5)
     print(f"  Threshold from config: {threshold}")
-    
-    # Tokenize
-    def tokenize(ex):
-        return tokenizer(ex["text"], truncation=True, max_length=2048, padding=False)
-    test_tok = test_ds.map(tokenize, batched=True)
-    test_tok.set_format("torch", columns=["input_ids", "attention_mask", "labels"])
-    
+
+    # Tokenize into plain lists (avoids set_format/slicing bugs)
+    texts = test_ds["text"]
+    labels_list = test_ds["labels"]
+    encodings = tokenizer(texts, truncation=True, max_length=2048, padding=True)
+    tensor_ds = TensorDataset(encodings, labels_list)
+    loader = DataLoader(tensor_ds, batch_size=32, shuffle=False)
+
     # Predict
     all_probs = []
     all_labels = []
-    bs = 32
-    for i in range(0, len(test_tok), bs):
-        batch = test_tok[i:i+bs]
-        inputs = {k: v.to(device) for k, v in batch.items() if k != "labels"}
-        with torch.no_grad():
-            logits = model(**inputs).logits
+    with torch.no_grad():
+        for batch in loader:
+            input_ids = batch["input_ids"].to(device)
+            attention_mask = batch["attention_mask"].to(device)
+            logits = model(input_ids=input_ids, attention_mask=attention_mask).logits
             probs = torch.softmax(logits, dim=-1).cpu().numpy()
-        all_probs.append(probs)
-        all_labels.extend(batch["labels"].cpu().numpy().tolist())
-    
+            all_probs.append(probs)
+            all_labels.extend(batch["label"].cpu().numpy().tolist())
+
     probs = np.vstack(all_probs)
     labels = np.array(all_labels)
-    
+
     # Default predictions
     preds_default = np.argmax(probs, axis=-1)
     acc_default = accuracy_score(labels, preds_default)
     f1_default = f1_score(labels, preds_default, average="macro", zero_division=0)
-    
-    # Calibrated predictions (binary tasks)
+
+    print(f"  Default:  acc={acc_default:.4f}, f1_macro={f1_default:.4f}")
+
     if num_labels == 2:
+        # Calibrated predictions
         preds_cal = (probs[:, 1] >= threshold).astype(int)
         acc_cal = accuracy_score(labels, preds_cal)
         f1_cal = f1_score(labels, preds_cal, average="macro", zero_division=0)
-        
-        print(f"  Default:  acc={acc_default:.4f}, f1_macro={f1_default:.4f}")
+
+        # Per-class precision/recall
+        p, r, f1, support = precision_recall_fscore_support(labels, preds_cal, zero_division=0)
+
         print(f"  Calibrated (t={threshold:.3f}): acc={acc_cal:.4f}, f1_macro={f1_cal:.4f}")
+
+        # DETECT COLLAPSE: if all predictions are same class
+        unique_preds = np.unique(preds_cal)
+        if len(unique_preds) == 1:
+            print(f"  ⚠️  MAJORITY-CLASS COLLAPSE DETECTED: model predicts only class {unique_preds[0]}")
+            print(f"     Accuracy = base rate of class {unique_preds[0]} = {max(pct=(labels==unique_preds[0]).mean()):.1%}")
+
         print(f"\n  Classification Report (calibrated):")
-        print(f"  {classification_report(labels, preds_cal, target_names=['neg','pos'], zero_division=0)}")
-        
-        return {"accuracy": acc_cal, "f1_macro": f1_cal, "threshold": threshold,
-                "accuracy_default": acc_default, "f1_default": f1_default}
+        print(f"  {classification_report(labels, preds_cal, target_names=['neg','pos'], zero_division=0, digits=4)}")
+        print(f"  Per-class: neg P={p[0]:.4f} R={r[0]:.4f} F1={f1[0]:.4f} | pos P={p[1]:.4f} R={r[1]:.4f} F1={f1[1]:.4f}")
+
+        return {
+            "accuracy": acc_cal, "f1_macro": f1_cal,
+            "accuracy_default": acc_default, "f1_default": f1_default,
+            "threshold": threshold,
+            "per_class": {
+                "neg": {"precision": float(p[0]), "recall": float(r[0]), "f1": float(f1[0]), "support": int(support[0])},
+                "pos": {"precision": float(p[1]), "recall": float(r[1]), "f1": float(f1[1]), "support": int(support[1])},
+            },
+            "collapsed": len(unique_preds) == 1,
+            "class_dist": lc,
+        }
     else:
-        print(f"  Accuracy: {acc_default:.4f}, F1 macro: {f1_default:.4f}")
+        # Multi-class
+        p, r, f1_per, support_per = precision_recall_fscore_support(labels, preds_default, zero_division=0)
         print(f"\n  Classification Report:")
-        print(f"  {classification_report(labels, preds_default, zero_division=0)}")
-        return {"accuracy": acc_default, "f1_macro": f1_default, "threshold": None}
+        print(f"  {classification_report(labels, preds_default, zero_division=0, digits=4)}")
+
+        per_class = {}
+        for i in range(num_labels):
+            per_class[str(i)] = {"precision": float(p[i]), "recall": float(r[i]), "f1": float(f1_per[i]), "support": int(support_per[i])}
+
+        return {
+            "accuracy": acc_default, "f1_macro": f1_default,
+            "threshold": None,
+            "per_class": per_class,
+            "collapsed": np.unique(preds_default).size == 1,
+            "class_dist": lc,
+        }
+
+# ═══════════════════════════════════════════
+# Main
+# ═══════════════════════════════════════════
 
 def main():
     results = {}
-    
-    for task_name in ["tier_router", "tool_gater", "verifier_gater"]:
+
+    for task_name in TASK_NAMES:
         num_labels = NUM_LABELS_MAP[task_name]
-        
+
         # Evaluate v2
+        print(f"\n{'#'*60}")
+        print(f"### V2 MODEL: {task_name}")
+        print(f"{'#'*60}")
         v2_res = evaluate_model(V2_MODELS[task_name], task_name, num_labels)
-        
+
         # Evaluate v1
+        print(f"\n{'#'*60}")
+        print(f"### V1 MODEL: {task_name} (baseline)")
+        print(f"{'#'*60}")
         v1_res = evaluate_model(V1_MODELS[task_name], task_name, num_labels)
-        
+
         if v2_res and v1_res:
-            delta = v2_res["f1_macro"] - v1_res["f1_macro"]
-            print(f"\n  >>> v1->v2 F1 delta: {v1_res['f1_macro']:.4f} -> {v2_res['f1_macro']:.4f} = {delta:+.4f}")
-        
-        results[task_name] = {"v1": v1_res, "v2": v2_res}
-    
+            delta_f1 = v2_res["f1_macro"] - v1_res["f1_macro"]
+            delta_acc = v2_res["accuracy"] - v1_res["accuracy"]
+            print(f"\n  >>> v1 → v2 delta: F1 {v1_res['f1_macro']:.4f} → {v2_res['f1_macro']:.4f} = {delta_f1:+.4f}")
+            print(f"  >>> v1 → v2 delta: Acc {v1_res['accuracy']:.4f} → {v2_res['accuracy']:.4f} = {delta_acc:+.4f}")
+            results[task_name] = {"v1": v1_res, "v2": v2_res, "delta_f1": delta_f1, "delta_acc": delta_acc}
+        else:
+            results[task_name] = {"v1": v1_res, "v2": v2_res, "error": "One or both models failed"}
+
+    # Final summary
     print(f"\n{'='*60}")
     print("FINAL COMPARISON")
     print(f"{'='*60}")
+
+    for tn in TASK_NAMES:
+        r = results.get(tn, {})
+        v1ok = r.get("v1") and not r["v1"].get("collapsed") if r.get("v1") else False
+        v2ok = r.get("v2") and not r["v2"].get("collapsed") if r.get("v2") else False
+        v1collapsed = r.get("v1", {}).get("collapsed", False)
+        v2collapsed = r.get("v2", {}).get("collapsed", False)
+        delta = r.get("delta_f1", float("nan"))
+
+        status = "OK"
+        if v2collapsed:
+            status = "⚠️  V2 COLLAPSED"
+        elif v1collapsed:
+            status = "⚠️  V1 COLLAPSED"
+
+        print(f"  {tn:<20} v1_f1={r.get('v1',{}).get('f1_macro',0):.4f}  "
+              f"v2_f1={r.get('v2',{}).get('f1_macro',0):.4f}  "
+              f"delta={delta:+.4f}  {status}")
+
     print(json.dumps(results, indent=2, default=str))
-    
+
     # Save results
     with open("/tmp/v2_verification_results.json", "w") as f:
         json.dump(results, f, indent=2, default=str)
-    
+
     # Push results
     from huggingface_hub import HfApi
     api = HfApi()
@@ -219,7 +317,7 @@ def main():
         repo_id="narcolepticchicken/agent-cost-optimizer",
         repo_type="model",
     )
-    print("\nResults pushed to agent-cost-optimizer repo")
+    print("\nResults pushed to narcolepticchicken/agent-cost-optimizer")
 
 if __name__ == "__main__":
     main()