"""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'(.*?)', 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()