Add synthetic data + full training pipeline

synthetic_data_and_train.py

@@ -0,0 +1,535 @@
+#!/usr/bin/env python3
+"""
+Speculative Tool Actions — Synthetic Data + Training Pipeline
+==============================================================
+Generates synthetic agent traces with 9 action types, trains proposer + verifier,
+evaluates all 5 configs, and produces cost-quality frontier report.
+
+Designed to run as a single HF Job on GPU hardware.
+"""
+import os, sys, json, re, random, math, subprocess, time
+from collections import Counter, defaultdict
+from datetime import datetime
+
+# Install required packages
+subprocess.check_call([sys.executable, "-m", "pip", "install", "--quiet",
+    "datasets", "transformers", "trl", "peft", "accelerate", "huggingface_hub", "trackio", "torch"])
+
+import torch
+from datasets import Dataset, DatasetDict
+from transformers import AutoModelForCausalLM, AutoTokenizer, set_seed
+from peft import LoraConfig, get_peft_model
+from trl import SFTTrainer, SFTConfig, RewardTrainer, RewardConfig
+
+set_seed(42)
+torch.manual_seed(42)
+random.seed(42)
+
+HUB_ORG = "narcolepticchicken"
+ACTION_TYPES = [
+    "tool_call", "retrieval", "file_read", "file_write",
+    "repair", "verifier", "ask_clarification", "final_answer", "BLOCKED",
+]
+
+COST = {"strong_in": 1.0, "strong_out": 1.0, "cheap_in": 0.2, "cheap_out": 0.2}
+
+# ========================================================================
+# Synthetic Data Generator
+# ========================================================================
+TASK_TEMPLATES = [
+    "Fix a bug in the authentication module.",
+    "Implement a new search feature.",
+    "Write unit tests for the API layer.",
+    "Refactor the database connection pool.",
+    "Add logging to the payment gateway.",
+    "Update documentation for the CLI tool.",
+    "Debug a memory leak in the worker process.",
+    "Optimize the image processing pipeline.",
+    "Integrate a third-party OAuth provider.",
+    "Set up CI/CD for the microservice.",
+    "Migrate from REST to GraphQL.",
+    "Add rate limiting to the public API.",
+    "Create a backup strategy for the database.",
+    "Audit the codebase for security vulnerabilities.",
+    "Implement caching for frequently accessed data.",
+]
+
+STATE_TEMPLATES = {
+    "tool_call": [
+        "I need to call the API to fetch user data.",
+        "Let me invoke the linter to check syntax.",
+        "I'll execute the test runner now.",
+        "Time to trigger the deployment script.",
+    ],
+    "retrieval": [
+        "I should search for similar issues in the tracker.",
+        "Let me look up the documentation for this function.",
+        "I'll query the knowledge base for best practices.",
+        "Need to find examples of this pattern online.",
+    ],
+    "file_read": [
+        "I need to read the configuration file first.",
+        "Let me check the existing implementation.",
+        "I'll examine the log file for clues.",
+        "Need to view the schema definition.",
+    ],
+    "file_write": [
+        "I'll write the fix to the source file.",
+        "Let me save the test cases to disk.",
+        "Need to update the requirements file.",
+        "I'll create a new migration script.",
+    ],
+    "repair": [
+        "The build is failing; let me fix the import error.",
+        "There's a null pointer exception to patch.",
+        "I need to correct the regex pattern.",
+        "Let me resolve the merge conflict.",
+    ],
+    "verifier": [
+        "Let me verify the fix by running tests.",
+        "I should check if the output is valid JSON.",
+        "Need to validate the schema changes.",
+        "I'll confirm the permissions are correct.",
+    ],
+    "ask_clarification": [
+        "I'm unsure about the expected behavior—could you clarify?",
+        "What is the target environment for this change?",
+        "Do you want me to preserve backward compatibility?",
+        "Which branch should I base this on?",
+    ],
+    "final_answer": [
+        "The task is complete. Summary of changes: ...",
+        "All tests pass. Here's the final solution.",
+        "Deployment successful. Verification complete.",
+        "Issue resolved. Closing the ticket.",
+    ],
+    "BLOCKED": [
+        "This request appears unsafe and I cannot proceed.",
+        "I'm sorry, but I cannot execute this command.",
+        "Blocked: the action violates safety policies.",
+        "Unsafe operation detected. Refusing to continue.",
+    ],
+}
+
+OBSERVATION_TEMPLATES = {
+    "tool_call": "Tool returned: status=200, data={...}",
+    "retrieval": "Found 3 relevant documents. Top result: ...",
+    "file_read": "File contents: 142 lines, class Foo { ... }",
+    "file_write": "File saved successfully. 3 lines changed.",
+    "repair": "Build passing. 0 errors, 2 warnings.",
+    "verifier": "Validation passed. Schema matches.",
+    "ask_clarification": "User replied: please use the main branch.",
+    "final_answer": "(no further action)",
+    "BLOCKED": "(no further action)",
+}
+
+
+def generate_trace(length=5, resolved_prob=0.8):
+    """Generate a single synthetic agent trace."""
+    task = random.choice(TASK_TEMPLATES)
+    messages = [{"role": "user", "content": task}]
+    gold_actions = []
+    for step in range(length):
+        # Choose action based on position in trace
+        if step == length - 1:
+            action = random.choices(
+                ["final_answer", "BLOCKED"],
+                weights=[0.85, 0.15]
+            )[0]
+        elif step == 0:
+            action = random.choices(
+                ["tool_call", "retrieval", "file_read", "ask_clarification"],
+                weights=[0.3, 0.25, 0.25, 0.2]
+            )[0]
+        else:
+            action = random.choice(ACTION_TYPES[:-2])  # exclude final_answer, BLOCKED
+
+        content = random.choice(STATE_TEMPLATES[action])
+        messages.append({"role": "assistant", "content": content})
+        gold_actions.append(action)
+        if action not in ("final_answer", "BLOCKED", "ask_clarification"):
+            messages.append({"role": "tool", "content": OBSERVATION_TEMPLATES[action]})
+
+    resolved = random.random() < resolved_prob
+    return messages, gold_actions, resolved
+
+
+def build_synthetic_datasets(n_train=5000, n_test=500):
+    print("=== Generating Synthetic Datasets ===")
+    p_rows, v_rows, e_rows = [], [], []
+
+    for _ in range(n_train + n_test):
+        msgs, actions, resolved = generate_trace(
+            length=random.randint(2, 6),
+            resolved_prob=0.75 if _ < n_train else 0.5
+        )
+        state = []
+        for i, msg in enumerate(msgs):
+            if msg["role"] == "assistant":
+                action = actions[len([m for m in msgs[:i] if m["role"] == "assistant"]) - 1]
+                comp = [{"role": "assistant", "content": msg["content"]}]
+                p_rows.append({"prompt": [m.copy() for m in state], "completion": comp, "action_type": action})
+                v_rows.append({"prompt": [m.copy() for m in state], "completion": comp, "label": resolved, "action_type": action})
+                e_rows.append({"messages": [m.copy() for m in state] + comp, "resolved": resolved, "action_type": action})
+            state.append(msg)
+
+    print(f"Total: proposer={len(p_rows)}, verifier={len(v_rows)}, eval={len(e_rows)}")
+    print("Distribution:", Counter(r["action_type"] for r in p_rows).most_common())
+
+    def fmt_proposer(r):
+        sys_msg = {"role": "system", "content": (
+            "You are an agent action predictor. Predict the next action from: "
+            + ", ".join(ACTION_TYPES) + ". Respond with exactly the action name.")}
+        prompt = [sys_msg] + r["prompt"]
+        if prompt:
+            prompt[-1]["content"] += "\n\n[Next Action] Choose one: " + ", ".join(ACTION_TYPES)
+        comp = r["completion"]
+        comp[0]["content"] = f"Action: {r['action_type']}\n" + comp[0]["content"]
+        return {"prompt": prompt, "completion": comp}
+
+    proposer_all = [fmt_proposer(r) for r in p_rows]
+    random.shuffle(proposer_all)
+    proposer_ds = DatasetDict({
+        "train": Dataset.from_list(proposer_all[:n_train]),
+        "test": Dataset.from_list(proposer_all[n_train:]),
+    })
+    proposer_ds.push_to_hub(f"{HUB_ORG}/speculative-actions-proposer-sft")
+    print("Pushed proposer dataset")
+
+    rng = random.Random(42)
+    good = [r for r in v_rows if r["label"]]
+    bad = [r for r in v_rows if not r["label"]]
+    if len(bad) < len(good) * 0.2:
+        for r in good:
+            wa = rng.choice([a for a in ACTION_TYPES if a != r["action_type"]])
+            bad.append({
+                "prompt": [m.copy() for m in r["prompt"]],
+                "completion": [{"role": "assistant", "content": f"Action: {wa}\n(incorrect action)"}],
+                "label": False, "action_type": wa,
+            })
+    pairs = []
+    for g in good:
+        b = rng.choice(bad)
+        pairs.append({
+            "prompt": [m.copy() for m in g["prompt"]],
+            "chosen": g["completion"],
+            "rejected": b["completion"],
+            "action_type": g["action_type"],
+        })
+    random.shuffle(pairs)
+    verifier_ds = DatasetDict({
+        "train": Dataset.from_list(pairs[:n_train]),
+        "test": Dataset.from_list(pairs[n_train:]),
+    })
+    verifier_ds.push_to_hub(f"{HUB_ORG}/speculative-actions-verifier-pref")
+    print("Pushed verifier dataset")
+
+    eval_all = e_rows
+    random.shuffle(eval_all)
+    eval_ds = Dataset.from_list(eval_all[:n_test])
+    eval_ds.push_to_hub(f"{HUB_ORG}/speculative-actions-eval")
+    print("Pushed eval dataset")
+
+    return proposer_ds, verifier_ds, eval_ds
+
+
+# ========================================================================
+# Training
+# ========================================================================
+def train_proposer():
+    print("\n=== Training Proposer ===")
+    ds = DatasetDict.load_from_disk(f"{HUB_ORG}/speculative-actions-proposer-sft") if False else None
+    # load from hub
+    from datasets import load_dataset
+    ds = load_dataset(f"{HUB_ORG}/speculative-actions-proposer-sft")
+
+    peft_config = LoraConfig(
+        r=16, lora_alpha=32,
+        target_modules=["q_proj", "v_proj", "k_proj", "o_proj"],
+        modules_to_save=["embed_tokens", "lm_head"],
+    )
+    config = SFTConfig(
+        output_dir="/tmp/proposer-out",
+        hub_model_id=f"{HUB_ORG}/speculative-proposer-qwen3-1.7b",
+        push_to_hub=True,
+        learning_rate=2e-4,
+        per_device_train_batch_size=4,
+        gradient_accumulation_steps=4,
+        num_train_epochs=2,
+        max_seq_length=2048,
+        bf16=True,
+        gradient_checkpointing=True,
+        logging_strategy="steps",
+        logging_steps=10,
+        logging_first_step=True,
+        disable_tqdm=True,
+        report_to="trackio",
+        run_name="proposer-sft-qwen3-1.7b",
+    )
+    trainer = SFTTrainer(
+        model="Qwen/Qwen3-1.7B",
+        train_dataset=ds["train"],
+        eval_dataset=ds["test"],
+        args=config,
+        peft_config=peft_config,
+    )
+    trainer.train()
+    trainer.push_to_hub()
+    print("Proposer training done.")
+
+
+def train_verifier():
+    print("\n=== Training Verifier ===")
+    from datasets import load_dataset
+    ds = load_dataset(f"{HUB_ORG}/speculative-actions-verifier-pref")
+
+    peft_config = LoraConfig(
+        r=16, lora_alpha=32,
+        target_modules=["q_proj", "v_proj", "k_proj", "o_proj"],
+        modules_to_save=["score"],
+    )
+    config = RewardConfig(
+        output_dir="/tmp/verifier-out",
+        hub_model_id=f"{HUB_ORG}/speculative-verifier-qwen3-4b",
+        push_to_hub=True,
+        learning_rate=1e-3,
+        per_device_train_batch_size=2,
+        gradient_accumulation_steps=8,
+        num_train_epochs=2,
+        max_seq_length=2048,
+        bf16=True,
+        gradient_checkpointing=True,
+        logging_strategy="steps",
+        logging_steps=10,
+        logging_first_step=True,
+        disable_tqdm=True,
+        report_to="trackio",
+        run_name="verifier-reward-qwen3-4b",
+    )
+    trainer = RewardTrainer(
+        model="Qwen/Qwen3-4B",
+        train_dataset=ds["train"],
+        eval_dataset=ds["test"],
+        args=config,
+        peft_config=peft_config,
+    )
+    trainer.train()
+    trainer.push_to_hub()
+    print("Verifier training done.")
+
+
+# ========================================================================
+# Evaluation
+# ========================================================================
+def parse_action(text):
+    for a in ACTION_TYPES:
+        if a.lower() in text.lower():
+            return a
+    return "tool_call"
+
+
+class EvalRunner:
+    def __init__(self, strong_name, cheap_name, verifier_name, device="cuda"):
+        self.device = device
+        print(f"Loading strong model: {strong_name}")
+        self.strong_tok = AutoTokenizer.from_pretrained(strong_name, trust_remote_code=True)
+        self.strong_model = AutoModelForCausalLM.from_pretrained(
+            strong_name, torch_dtype=torch.bfloat16, device_map="auto", trust_remote_code=True)
+        print(f"Loading cheap model: {cheap_name}")
+        self.cheap_tok = AutoTokenizer.from_pretrained(cheap_name, trust_remote_code=True)
+        self.cheap_model = AutoModelForCausalLM.from_pretrained(
+            cheap_name, torch_dtype=torch.bfloat16, device_map="auto", trust_remote_code=True)
+        self.verifier_name = verifier_name
+        if verifier_name:
+            print(f"Loading verifier: {verifier_name}")
+            self.v_tok = AutoTokenizer.from_pretrained(verifier_name, trust_remote_code=True)
+            self.v_model = AutoModelForCausalLM.from_pretrained(
+                verifier_name, torch_dtype=torch.bfloat16, device_map="auto", trust_remote_code=True)
+
+    def _gen(self, model, tokenizer, messages, max_new=64, temp=0.0):
+        inputs = tokenizer.apply_chat_template(
+            messages, tokenize=True, return_tensors="pt", add_generation_prompt=True
+        ).to(model.device)
+        with torch.no_grad():
+            out = model.generate(
+                inputs, max_new_tokens=max_new, do_sample=temp > 0,
+                temperature=temp if temp > 0 else None,
+                pad_token_id=tokenizer.pad_token_id or tokenizer.eos_token_id,
+            )
+        text = tokenizer.decode(out[0][inputs.shape[1]:], skip_special_tokens=True)
+        return text, inputs.shape[1], out.shape[1] - inputs.shape[1]
+
+    def run_a(self, messages):
+        s = {"role": "system", "content": f"Predict next action from: {', '.join(ACTION_TYPES)}"}
+        out, i, o = self._gen(self.strong_model, self.strong_tok, [s] + messages)
+        return parse_action(out), i, o, "strong"
+
+    def run_b(self, messages):
+        s = {"role": "system", "content": f"Predict next action from: {', '.join(ACTION_TYPES)}"}
+        out, i, o = self._gen(self.cheap_model, self.cheap_tok, [s] + messages)
+        return parse_action(out), i, o, "cheap"
+
+    def run_c(self, messages):
+        s = {"role": "system", "content": f"Predict next action from: {', '.join(ACTION_TYPES)}"}
+        prop, i1, o1 = self._gen(self.cheap_model, self.cheap_tok, [s] + messages)
+        vp = messages + [{"role": "assistant", "content": prop},
+                         {"role": "user", "content": "Is this action correct? Answer ONLY yes or no."}]
+        ver, i2, o2 = self._gen(self.strong_model, self.strong_tok, vp, max_new=10)
+        if "yes" in ver.lower():
+            return parse_action(prop), i1 + i2, o1 + o2, "mixed"
+        out, i3, o3 = self._gen(self.strong_model, self.strong_tok, [s] + messages)
+        return parse_action(out), i1 + i2 + i3, o1 + o2 + o3, "mixed"
+
+    def run_d(self, messages):
+        if not self.verifier_name:
+            raise ValueError("Need verifier")
+        s = {"role": "system", "content": f"Predict next action from: {', '.join(ACTION_TYPES)}"}
+        prop, i1, o1 = self._gen(self.cheap_model, self.cheap_tok, [s] + messages)
+        vp = messages + [{"role": "assistant", "content": prop},
+                         {"role": "user", "content": "Rate this action: good or bad."}]
+        ver, i2, o2 = self._gen(self.v_model, self.v_tok, vp, max_new=10)
+        if "good" in ver.lower():
+            return parse_action(prop), i1 + i2, o1 + o2, "cheap"
+        out, i3, o3 = self._gen(self.strong_model, self.strong_tok, [s] + messages)
+        return parse_action(out), i1 + i2 + i3, o1 + o2 + o3, "mixed"
+
+    def run_e(self, messages, n=3):
+        s = {"role": "system", "content": f"Predict next action from: {', '.join(ACTION_TYPES)}"}
+        props = []
+        ti, to = 0, 0
+        for _ in range(n):
+            p, i, o = self._gen(self.cheap_model, self.cheap_tok, [s] + messages, temp=0.7)
+            props.append(p); ti += i; to += o
+        best = props[0]; best_score = -1
+        for p in props:
+            rp = messages + [{"role": "assistant", "content": p},
+                             {"role": "user", "content": "Score 1-10."}]
+            st, i, o = self._gen(self.strong_model, self.strong_tok, rp, max_new=5)
+            ti += i; to += o
+            m = re.search(r'(\d+)', st)
+            if m:
+                sc = int(m.group(1))
+                if sc > best_score:
+                    best_score = sc; best = p
+        return parse_action(best), ti, to, "mixed"
+
+
+def evaluate(limit=200):
+    print("\n=== Evaluation ===")
+    from datasets import load_dataset
+    ds = load_dataset(f"{HUB_ORG}/speculative-actions-eval", split="train")
+    ds = ds.shuffle(seed=42).select(range(min(limit, len(ds))))
+
+    runner = EvalRunner(
+        strong_name="Qwen/Qwen2.5-7B-Instruct",
+        cheap_name="Qwen/Qwen3-1.7B",
+        verifier_name=f"{HUB_ORG}/speculative-verifier-qwen3-4b",
+    )
+
+    results = defaultdict(lambda: {"correct": 0, "total": 0, "cost": 0.0, "unsafe": 0})
+    for idx, ex in enumerate(ds):
+        msgs = ex["messages"]; gold = ex["action_type"]
+        for cfg, func in [("A", runner.run_a), ("B", runner.run_b),
+                          ("C", runner.run_c), ("D", runner.run_d),
+                          ("E", lambda m: runner.run_e(m, n=3))]:
+            try:
+                pred, i_t, o_t, mtype = func(msgs)
+            except Exception as e:
+                print(f"Error {cfg} idx {idx}: {e}")
+                pred, i_t, o_t, mtype = "tool_call", 0, 0, "unknown"
+            results[cfg]["total"] += 1
+            if pred == gold:
+                results[cfg]["correct"] += 1
+            if pred == "BLOCKED" and gold != "BLOCKED":
+                results[cfg]["unsafe"] += 1
+            if pred != "BLOCKED" and gold == "BLOCKED":
+                results[cfg]["unsafe"] += 1
+            results[cfg]["cost"] += i_t * COST.get(f"{mtype}_in", 1.0) + o_t * COST.get(f"{mtype}_out", 1.0)
+        if (idx + 1) % 20 == 0:
+            print(f"  Evaluated {idx + 1}/{min(limit, len(ds))}")
+
+    for cfg in results:
+        t = max(results[cfg]["total"], 1)
+        results[cfg]["accuracy"] = results[cfg]["correct"] / t
+        results[cfg]["avg_cost"] = results[cfg]["cost"] / t
+        results[cfg]["unsafe_rate"] = results[cfg]["unsafe"] / t
+
+    summary = {k: dict(v) for k, v in results.items()}
+    with open("/tmp/eval_results.json", "w") as f:
+        json.dump(summary, f, indent=2)
+    print(json.dumps(summary, indent=2))
+    return summary
+
+
+# ========================================================================
+# Report
+# ========================================================================
+def generate_report(eval_results):
+    print("\n=== Generating Report ===")
+    lines = ["# Speculative Tool Actions — Ablation Report\n\n"]
+    lines.append("## Configurations\n\n")
+    lines.append("- **A**: Always strong model (Qwen2.5-7B)\n")
+    lines.append("- **B**: Cheap model only (Qwen3-1.7B)\n")
+    lines.append("- **C**: Cheap proposer + strong verifier\n")
+    lines.append("- **D**: Cheap proposer + trained trace judge (Qwen3-4B reward model)\n")
+    lines.append("- **E**: Multi-proposal reranking (3 cheap proposals + strong scoring)\n\n")
+
+    lines.append("## Results\n\n")
+    lines.append("| Config | Accuracy | Avg Cost | Unsafe-Action Rate |\n")
+    lines.append("|--------|----------|----------|-------------------|\n")
+    for cfg in sorted(eval_results):
+        r = eval_results[cfg]
+        lines.append(f"| {cfg} | {r['accuracy']:.3f} | {r['avg_cost']:.2f} | {r['unsafe_rate']:.3f} |\n")
+
+    lines.append("\n## Cost-Quality Frontier\n\n")
+    points = [(r["avg_cost"], r["accuracy"], cfg) for cfg, r in eval_results.items()]
+    points.sort()
+    frontier = []
+    max_acc = -1
+    for cost, acc, cfg in points:
+        if acc > max_acc:
+            frontier.append((cost, acc, cfg)); max_acc = acc
+    lines.append("Pareto-optimal configs:\n")
+    for cost, acc, cfg in frontier:
+        lines.append(f"- **{cfg}**: cost={cost:.2f}, accuracy={acc:.3f}\n")
+
+    lines.append("\n## Recommendations\n\n")
+    best_ratio = None; best_cfg = None
+    for cfg, r in eval_results.items():
+        ratio = r["accuracy"] / max(r["avg_cost"], 0.01)
+        if best_ratio is None or ratio > best_ratio:
+            best_ratio = ratio; best_cfg = cfg
+    lines.append(f"- **Best accuracy/cost ratio**: Config {best_cfg} (ratio={best_ratio:.3f})\n")
+
+    best_acc_cfg = max(eval_results, key=lambda c: eval_results[c]["accuracy"])
+    lines.append(f"- **Highest accuracy**: Config {best_acc_cfg} ({eval_results[best_acc_cfg]['accuracy']:.3f})\n")
+
+    best_acc = eval_results[best_acc_cfg]["accuracy"]
+    threshold = best_acc * 0.9
+    cheap = {c: r for c, r in eval_results.items() if r["accuracy"] >= threshold}
+    if cheap:
+        cheapest = min(cheap, key=lambda c: cheap[c]["avg_cost"])
+        lines.append(f"- **Cheapest within 90% of best accuracy**: Config {cheapest} "
+                     f"(cost={cheap[cheapest]['avg_cost']:.2f}, acc={cheap[cheapest]['accuracy']:.3f})\n")
+
+    report = "".join(lines)
+    with open("/tmp/ablation_report.md", "w") as f:
+        f.write(report)
+    print(report)
+    return report
+
+
+# ========================================================================
+# Main
+# ========================================================================
+def main():
+    build_synthetic_datasets(n_train=5000, n_test=500)
+    train_proposer()
+    train_verifier()
+    results = evaluate(limit=200)
+    generate_report(results)
+    print("\n=== Pipeline Complete ===")
+
+
+if __name__ == "__main__":
+    main()