Add standalone dataset build job script

build_datasets_job.py

@@ -0,0 +1,163 @@
+"""
+HF Jobs script: Build Speculative Tool Actions datasets and push to Hub.
+"""
+import json
+import re
+from collections import Counter
+from datasets import load_dataset, Dataset
+from random import Random
+
+ACTION_TYPES = [
+    "tool_call", "retrieval", "file_read", "file_write",
+    "repair", "verifier", "ask_clarification", "final_answer", "BLOCKED",
+]
+
+
+def classify_action(content, tool_calls=None):
+    c = content.lower()
+    tc = json.dumps(tool_calls).lower() if tool_calls else ""
+    combined = c + " " + tc
+    if re.search(r'\b(final answer|conclusion|summary:|in conclusion|the answer is)\b', combined):
+        return "final_answer"
+    if re.search(r'\b(ask for clarification|need more info|could you clarify|what do you mean)\b', combined):
+        return "ask_clarification"
+    if re.search(r'\b(blocked|unsafe|i cannot|i\'m sorry, but|refuse|not allowed|harmful)\b', combined):
+        return "BLOCKED"
+    if re.search(r'\b(write.*file|save.*file|edit.*file|patch|diff)\b', combined):
+        return "file_write"
+    if re.search(r'\b(read.*file|view.*file|cat |head |tail |open.*file|get_content)\b', combined):
+        return "file_read"
+    if re.search(r'\b(repair|fix.*bug|correct.*error|debug|resolve|try.*again with)\b', combined):
+        return "repair"
+    if re.search(r'\b(verify|check|validate|test|assert|review)\b', combined):
+        return "verifier"
+    if re.search(r'\b(search|retrieve|find|lookup|query|google|bing)\b', combined):
+        return "retrieval"
+    if tool_calls or re.search(r'\b(function call|tool call|invoke|execute)\b', combined):
+        return "tool_call"
+    return "tool_call"
+
+
+def process_swe_smith(max_rows=5000):
+    print("Loading SWE-smith tool/train ...")
+    ds = load_dataset("SWE-bench/SWE-smith-trajectories", "tool", split="train", streaming=True)
+    rows_proposer, rows_verifier, rows_eval = [], [], []
+    count = 0
+    for example in ds:
+        count += 1
+        if count > max_rows:
+            break
+        messages = example.get("messages", [])
+        resolved = example.get("resolved", False)
+        state_so_far = []
+        for msg in messages:
+            role = msg.get("role", "")
+            content = msg.get("content", "")
+            tool_calls = msg.get("tool_calls", None)
+            if role in ("assistant", "agent"):
+                action_type = classify_action(content, tool_calls)
+                prompt_messages = state_so_far.copy()
+                completion_messages = [{"role": "assistant", "content": content}]
+                if tool_calls:
+                    completion_messages[0]["tool_calls"] = tool_calls
+                rows_proposer.append({"prompt": prompt_messages, "completion": completion_messages, "action_type": action_type})
+                rows_verifier.append({"prompt": prompt_messages, "completion": completion_messages, "label": bool(resolved), "action_type": action_type})
+                rows_eval.append({"messages": prompt_messages + completion_messages, "resolved": resolved, "action_type": action_type})
+            state_so_far.append(msg)
+    print(f"  -> proposer={len(rows_proposer)}, verifier={len(rows_verifier)}")
+    return rows_proposer, rows_verifier, rows_eval
+
+
+def process_toolbench(max_rows=3000):
+    print("Loading toolbench/train ...")
+    ds = load_dataset("tuandunghcmut/toolbench-v1", split="train", streaming=True)
+    rows_proposer, rows_verifier, rows_eval = [], [], []
+    count = 0
+    for example in ds:
+        count += 1
+        if count > max_rows:
+            break
+        conv = example.get("conversations", {})
+        froms = conv.get("from", [])
+        values = conv.get("value", [])
+        state_so_far = []
+        for role, content in zip(froms, values):
+            msg = {"role": role, "content": content}
+            if role == "assistant":
+                action_type = classify_action(content)
+                rows_proposer.append({"prompt": state_so_far.copy(), "completion": [msg], "action_type": action_type})
+                rows_verifier.append({"prompt": state_so_far.copy(), "completion": [msg], "label": True, "action_type": action_type})
+                rows_eval.append({"messages": state_so_far + [msg], "resolved": True, "action_type": action_type})
+            state_so_far.append(msg)
+    print(f"  -> proposer={len(rows_proposer)}, verifier={len(rows_verifier)}")
+    return rows_proposer, rows_verifier, rows_eval
+
+
+def build_proposer(rows, hub_org):
+    def fmt(row):
+        system_msg = {"role": "system", "content": (
+            "You are an agent action predictor. Given the conversation state, predict the next action from: "
+            + ", ".join(ACTION_TYPES) + ". Respond with exactly the action name and a brief justification.")}
+        prompt = [system_msg] + row["prompt"]
+        prompt[-1]["content"] += "\n\n[Next Action Prediction] Choose one: " + ", ".join(ACTION_TYPES)
+        completion = row["completion"]
+        completion[0]["content"] = f"Action: {row['action_type']}\n" + completion[0]["content"]
+        return {"prompt": prompt, "completion": completion}
+    ds = Dataset.from_list([fmt(r) for r in rows])
+    ds = ds.shuffle(seed=42).train_test_split(test_size=0.1)
+    ds.push_to_hub(f"{hub_org}/speculative-actions-proposer-sft")
+    print(f"Pushed proposer SFT to {hub_org}/speculative-actions-proposer-sft")
+    return ds
+
+
+def build_verifier(rows, hub_org):
+    rng = Random(42)
+    good_rows = [r for r in rows if r["label"]]
+    bad_rows = [r for r in rows if not r["label"]]
+    if len(bad_rows) < len(good_rows) * 0.2:
+        for r in good_rows:
+            wrong_action = rng.choice([a for a in ACTION_TYPES if a != r["action_type"]])
+            bad_rows.append({
+                "prompt": r["prompt"],
+                "completion": [{"role": "assistant", "content": f"Action: {wrong_action}\n(synthetic incorrect action)"}],
+                "label": False,
+                "action_type": wrong_action,
+            })
+    pairs = []
+    for g in good_rows:
+        b = rng.choice(bad_rows)
+        pairs.append({"prompt": g["prompt"], "chosen": g["completion"], "rejected": b["completion"], "action_type": g["action_type"]})
+    ds = Dataset.from_list(pairs)
+    ds = ds.shuffle(seed=42).train_test_split(test_size=0.1)
+    ds.push_to_hub(f"{hub_org}/speculative-actions-verifier-pref")
+    print(f"Pushed verifier pref to {hub_org}/speculative-actions-verifier-pref")
+    return ds
+
+
+def build_eval(rows, hub_org):
+    ds = Dataset.from_list(rows)
+    ds = ds.shuffle(seed=42).select(range(min(2000, len(rows))))
+    ds.push_to_hub(f"{hub_org}/speculative-actions-eval")
+    print(f"Pushed eval to {hub_org}/speculative-actions-eval")
+    return ds
+
+
+def main():
+    hub_org = "narcolepticchicken"
+    p1, v1, e1 = process_swe_smith(5000)
+    p2, v2, e2 = process_toolbench(3000)
+    proposer_rows = p1 + p2
+    verifier_rows = v1 + v2
+    eval_rows = e1 + e2
+    print(f"\nTotal: proposer={len(proposer_rows)}, verifier={len(verifier_rows)}, eval={len(eval_rows)}")
+    print("\nAction distribution:")
+    for act, n in Counter(r["action_type"] for r in proposer_rows).most_common():
+        print(f"  {act}: {n}")
+    build_proposer(proposer_rows, hub_org)
+    build_verifier(verifier_rows, hub_org)
+    build_eval(eval_rows, hub_org)
+    print("Done.")
+
+
+if __name__ == "__main__":
+    main()