Upload build_traces.py

build_traces.py

@@ -0,0 +1,241 @@
+"""Build and publish the ACO training traces dataset.
+
+Loads the same three source datasets used for training:
+- lockon/ToolACE → tool_gater traces (query→tool_called binary)
+- RouteWorks/RouterArena → tier_router traces (question→difficulty tier)
+- R2E-Gym/R2EGym-Verifier-Trajectories → verifier_gater traces (patch→verified binary)
+
+Applies identical preprocessing to verify_v2.py training. 
+Publishes train/test splits to narcolepticchicken/aco-traces.
+
+Usage:
+    uv run --with transformers,torch,datasets,huggingface_hub build_traces.py
+"""
+import json
+import os
+import re
+from datasets import Dataset, load_dataset
+from huggingface_hub import HfApi
+
+# ═══════════════════════════════════════════════════════════════
+# 1. Tool Gater: ToolACE → binary tool-call classification
+# ═══════════════════════════════════════════════════════════════
+
+def build_tool_gater():
+    ds = load_dataset("lockon/ToolACE", split="train")
+    texts, labels = [], []
+    skipped = 0
+    for row in ds:
+        conv = row.get("conversations", [])
+        q = ""
+        for turn in conv:
+            if turn.get("from") == "user":
+                q = turn.get("value", "")[:1500]
+                break
+        if not q:
+            skipped += 1
+            continue
+        # Detect if any assistant turn includes a tool-call pattern
+        called = any(
+            re.search(r'\[[A-Z][a-zA-Z]+\s*\(', turn["value"])
+            for turn in conv if turn.get("from") == "assistant"
+        )
+        text = f"Query: {q}"
+        if row.get("system"):
+            text = f"System: {row['system'][:500]}\n\n{text}"
+        texts.append(text[:2000])
+        labels.append(1 if called else 0)
+    ds = Dataset.from_dict({"text": texts, "labels": labels})
+    print(f"  ToolACE: {len(texts)} samples, {skipped} skipped, "
+          f"pos={sum(labels)} ({sum(labels)/len(labels)*100:.1f}%)")
+    return ds.train_test_split(test_size=0.15, seed=42)
+
+# ═══════════════════════════════════════════════════════════════
+# 2. Tier Router: RouterArena → 3-class difficulty classification
+# ═══════════════════════════════════════════════════════════════
+
+def build_tier_router():
+    ds = load_dataset("RouteWorks/RouterArena", "default", split="full")
+    tmap = {"easy": 0, "medium": 1, "hard": 2}
+    texts, labels = [], []
+    skipped = 0
+    for row in ds:
+        d = row.get("Difficulty", "").strip().lower()
+        if d not in tmap:
+            skipped += 1
+            continue
+        parts = []
+        if row.get("Domain"):
+            parts.append(f"[{row['Domain']}]")
+        if row.get("Context"):
+            parts.append(f"Context: {row['Context']}")
+        parts.append(row.get("Question", ""))
+        o = row.get("Options", "")
+        if o:
+            parts.append(f"Options: {'; '.join(o) if isinstance(o, list) else o}")
+        texts.append(" ".join(parts)[:2000])
+        labels.append(tmap[d])
+    ds = Dataset.from_dict({"text": texts, "labels": labels})
+    dist = {k: labels.count(k) for k in [0,1,2]}
+    print(f"  RouterArena: {len(texts)} samples, {skipped} skipped, "
+          f"dist={dist}")
+    return ds.train_test_split(test_size=0.15, seed=42)
+
+# ═══════════════════════════════════════════════════════════════
+# 3. Verifier Gater: R2E-Gym → binary verification classification
+# ═══════════════════════════════════════════════════════════════
+
+def build_verifier_gater():
+    ds = load_dataset("R2E-Gym/R2EGym-Verifier-Trajectories", split="train")
+    texts, labels = [], []
+    for row in ds:
+        messages = row["messages"]
+        fl = messages[1]["content"] if len(messages) > 1 else ""
+
+        # Extract github issue text
+        task_text = ""
+        for msg in messages:
+            if msg["role"] == "user" and "INTERACTION LOG" in msg["content"]:
+                m = re.search(r'<github_issue>(.*?)</github_issue>', msg["content"], re.DOTALL)
+                if m:
+                    task_text = m.group(1).strip()[:1000]
+                    break
+        if not task_text:
+            for msg in messages:
+                if msg["role"] == "system":
+                    task_text = msg["content"][:500]
+                    break
+
+        # Extract agent action summary from last 3 turns
+        ab = re.findall(r'\[ASSISTANT\](.*?)(?:\[USER\]|\[STEP\]|$)', fl, re.DOTALL)
+        agent_sum = " ".join(b.strip()[:200] for b in ab[-3:])
+
+        # Extract patch
+        pm = re.search(r'=== FINAL PATCH ===\s*\n(.*?)\n=== END FINAL PATCH ===', fl, re.DOTALL)
+        patch = pm.group(1)[:500] if pm else ""
+
+        text = f"TASK: {task_text[:600]}\nAGENT_ACTIONS: {agent_sum[:600]}\nPATCH: {patch[:400]}"
+        texts.append(text[:2000])
+        labels.append(1 if row["rewards"] >= 1.0 else 0)
+
+    ds = Dataset.from_dict({"text": texts, "labels": labels})
+    print(f"  R2E-Gym: {len(texts)} samples, pos={sum(labels)} ({sum(labels)/len(labels)*100:.1f}%)")
+    return ds.train_test_split(test_size=0.15, seed=42)
+
+# ═══════════════════════════════════════════════════════════════
+# Build and publish
+# ═══════════════════════════════════════════════════════════════
+
+def main():
+    api = HfApi()
+    repo = "narcolepticchicken/aco-traces"
+
+    builders = {
+        "tool_gater": build_tool_gater,
+        "tier_router": build_tier_router,
+        "verifier_gater": build_verifier_gater,
+    }
+
+    metadata = {
+        "description": "Agent Cost Optimizer training traces. Preprocessed from "
+                       "lockon/ToolACE, RouteWorks/RouterArena, and "
+                       "R2E-Gym/R2EGym-Verifier-Trajectories.",
+        "license": "apache-2.0",
+        "citation_sources": [
+            "ToolACE (arXiv:2409.00920)",
+            "RouteWorks RouterArena (arXiv:2510.00202)",
+            "R2E-Gym Verifier Trajectories",
+        ],
+        "build_date": None,  # will be set
+        "preprocessing": "Identical to verify_v2.py loaders. "
+                         "Tool gater: regex tool-call detection in ToolACE conversations. "
+                         "Tier router: RouterArena Difficulty field → 3 classes. "
+                         "Verifier gater: R2E-Gym rewards threshold 1.0, "
+                         "600-char github issue + agent action summary + patch snippet.",
+    }
+
+    from datetime import datetime
+    metadata["build_date"] = datetime.utcnow().isoformat()
+
+    for task_name, builder_fn in builders.items():
+        print(f"\n{'='*60}")
+        print(f"Building: {task_name}")
+        print(f"{'='*60}")
+
+        splits = builder_fn()
+        for split_name in ["train", "test"]:
+            ds = splits[split_name]
+            path = f"data/{task_name}/{split_name}.parquet"
+            ds.to_parquet(f"/tmp/{task_name}_{split_name}.parquet")
+
+            api.upload_file(
+                path_or_fileobj=f"/tmp/{task_name}_{split_name}.parquet",
+                path_in_repo=path,
+                repo_id=repo,
+                repo_type="dataset",
+            )
+            print(f"  Uploaded: {path} ({len(ds)} rows)")
+
+    # Upload metadata / data card
+    api.upload_file(
+        path_or_fileobj=json.dumps(metadata, indent=2).encode(),
+        path_in_repo="metadata.json",
+        repo_id=repo,
+        repo_type="dataset",
+    )
+
+    # Upload README
+    readme = f"""---
+license: apache-2.0
+task_categories:
+- text-classification
+language:
+- en
+tags:
+- agent-traces
+- cost-optimization
+- model-routing
+- tool-calling
+- verifier
+pretty_name: ACO Training Traces
+---
+
+# ACO Training Traces
+
+Training data for the Agent Cost Optimizer's specialist classifiers.
+
+## Source Datasets
+
+| Split | Source | Task | Classes |
+|-------|--------|------|---------|
+| tool_gater | lockon/ToolACE | Predict whether a tool call is needed | binary: no_tool (0) / call_tool (1) |
+| tier_router | RouteWorks/RouterArena | Predict difficulty tier | 3-class: easy (0) / medium (1) / hard (2) |
+| verifier_gater | R2E-Gym/R2EGym-Verifier-Trajectories | Predict whether patch passes verification | binary: fail (0) / pass (1) |
+
+## Preprocessing
+
+Identical to the loaders in `verify_v2.py`. See `metadata.json` for details.
+
+## Usage
+
+```python
+from datasets import load_dataset
+
+ds = load_dataset("narcolepticchicken/aco-traces", "tool_gater")
+# Contains 'train' and 'test' splits with 'text' and 'labels' columns
+```
+
+Built: {metadata['build_date']}
+"""
+    api.upload_file(
+        path_or_fileobj=readme.encode(),
+        path_in_repo="README.md",
+        repo_id=repo,
+        repo_type="dataset",
+    )
+
+    print(f"\nDataset published to https://huggingface.co/datasets/{repo}")
+    print(f"Total splits: {len(builders) * 2}")
+
+if __name__ == "__main__":
+    main()