Add end-to-end pipeline script

end_to_end_pipeline.py

@@ -0,0 +1,492 @@
+"""
+Speculative Tool Actions — End-to-End Pipeline
+================================================
+1. Build datasets from SWE-smith + ToolBench
+2. Train cheap proposer (Qwen3-1.7B + LoRA SFT)
+3. Train verifier (Qwen3-4B + LoRA Reward)
+4. Evaluate all 5 configs (A-E) on held-out set
+5. Generate ablation report + cost-quality frontier
+
+Run via: hf_jobs with GPU hardware (a10g-large or a100-large)
+"""
+import os
+import json
+import re
+import argparse
+from collections import Counter, defaultdict
+from random import Random
+
+import torch
+from datasets import load_dataset, Dataset
+from transformers import AutoModelForCausalLM, AutoTokenizer, set_seed
+from peft import PeftModel
+
+# ============================================================================
+# Configuration
+# ============================================================================
+HUB_ORG = "narcolepticchicken"
+ACTION_TYPES = [
+    "tool_call", "retrieval", "file_read", "file_write",
+    "repair", "verifier", "ask_clarification", "final_answer", "BLOCKED",
+]
+
+# Cost weights (relative)
+COST = {"strong_in": 1.0, "strong_out": 1.0, "cheap_in": 0.2, "cheap_out": 0.2}
+
+# ============================================================================
+# Step 1: Dataset Builder
+# ============================================================================
+def classify_action(content, tool_calls=None):
+    c = (content or "").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 build_datasets(max_swe=5000, max_toolbench=3000):
+    print("=== Step 1: Building Datasets ===")
+    # SWE-smith
+    ds_swe = load_dataset("SWE-bench/SWE-smith-trajectories", "tool", split="train", streaming=True)
+    p_rows, v_rows, e_rows = [], [], []
+    count = 0
+    for ex in ds_swe:
+        count += 1
+        if count > max_swe:
+            break
+        msgs = ex.get("messages", [])
+        resolved = ex.get("resolved", False)
+        state = []
+        for msg in msgs:
+            role = msg.get("role", "")
+            if role in ("assistant", "agent"):
+                atype = classify_action(msg.get("content", ""), msg.get("tool_calls"))
+                comp = [{"role": "assistant", "content": msg.get("content", "")}]
+                if msg.get("tool_calls"):
+                    comp[0]["tool_calls"] = msg["tool_calls"]
+                p_rows.append({"prompt": list(state), "completion": comp, "action_type": atype})
+                v_rows.append({"prompt": list(state), "completion": comp, "label": bool(resolved), "action_type": atype})
+                e_rows.append({"messages": list(state) + comp, "resolved": resolved, "action_type": atype})
+            state.append(msg)
+
+    # ToolBench
+    ds_tb = load_dataset("tuandunghcmut/toolbench-v1", split="train", streaming=True)
+    count = 0
+    for ex in ds_tb:
+        count += 1
+        if count > max_toolbench:
+            break
+        conv = ex.get("conversations", {})
+        for role, content in zip(conv.get("from", []), conv.get("value", [])):
+            msg = {"role": role, "content": content}
+            if role == "assistant":
+                atype = classify_action(content)
+                p_rows.append({"prompt": [m for m in state if m["role"] != "assistant" or m is not msg], "completion": [msg], "action_type": atype})
+                v_rows.append({"prompt": [m for m in state if m["role"] != "assistant" or m is not msg], "completion": [msg], "label": True, "action_type": atype})
+                e_rows.append({"messages": state + [msg], "resolved": True, "action_type": atype})
+            state.append(msg)
+
+    print(f"Total rows: proposer={len(p_rows)}, verifier={len(v_rows)}, eval={len(e_rows)}")
+    print("Action distribution:", Counter(r["action_type"] for r in p_rows).most_common())
+
+    # Proposer SFT dataset
+    def fmt_proposer(r):
+        sys_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 = [sys_msg] + r["prompt"]
+        if prompt:
+            prompt[-1]["content"] += "\n\n[Next Action Prediction] 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_ds = Dataset.from_list([fmt_proposer(r) for r in p_rows]).shuffle(seed=42).train_test_split(test_size=0.1)
+    proposer_ds.push_to_hub(f"{HUB_ORG}/speculative-actions-proposer-sft")
+    print(f"Pushed proposer dataset to {HUB_ORG}/speculative-actions-proposer-sft")
+
+    # Verifier preference dataset
+    rng = 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": r["prompt"],
+                "completion": [{"role": "assistant", "content": f"Action: {wa}\n(synthetic incorrect action)"}],
+                "label": False, "action_type": wa,
+            })
+    pairs = []
+    for g in good:
+        b = rng.choice(bad)
+        pairs.append({"prompt": g["prompt"], "chosen": g["completion"], "rejected": b["completion"], "action_type": g["action_type"]})
+    verifier_ds = Dataset.from_list(pairs).shuffle(seed=42).train_test_split(test_size=0.1)
+    verifier_ds.push_to_hub(f"{HUB_ORG}/speculative-actions-verifier-pref")
+    print(f"Pushed verifier dataset to {HUB_ORG}/speculative-actions-verifier-pref")
+
+    # Eval dataset
+    eval_ds = Dataset.from_list(e_rows).shuffle(seed=42).select(range(min(2000, len(e_rows))))
+    eval_ds.push_to_hub(f"{HUB_ORG}/speculative-actions-eval")
+    print(f"Pushed eval dataset to {HUB_ORG}/speculative-actions-eval")
+
+    return proposer_ds, verifier_ds, eval_ds
+
+
+# ============================================================================
+# Step 2: Train Proposer
+# ============================================================================
+def train_proposer():
+    print("\n=== Step 2: Training Proposer ===")
+    from trl import SFTTrainer, SFTConfig
+    from peft import LoraConfig
+
+    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=3,
+        max_seq_length=4096,
+        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 complete.")
+
+
+# ============================================================================
+# Step 3: Train Verifier
+# ============================================================================
+def train_verifier():
+    print("\n=== Step 3: Training Verifier ===")
+    from trl import RewardTrainer, RewardConfig
+    from peft import LoraConfig
+
+    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=4096,
+        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 complete.")
+
+
+# ============================================================================
+# Step 4: Evaluate Configs A-E
+# ============================================================================
+class EvalRunner:
+    def __init__(self, strong_name="Qwen/Qwen2.5-7B-Instruct", cheap_name="Qwen/Qwen3-1.7B",
+                 verifier_name=None, device="cuda"):
+        self.device = device
+        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)
+
+        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:
+            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=128, 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 _parse(self, text):
+        for a in ACTION_TYPES:
+            if a.lower() in text.lower():
+                return a
+        return "tool_call"
+
+    def run_a(self, messages):
+        sys_msg = {"role": "system", "content": f"Predict next action from: {', '.join(ACTION_TYPES)}"}
+        out, i_t, o_t = self._gen(self.strong_model, self.strong_tok, [sys_msg] + messages)
+        return self._parse(out), i_t, o_t, "strong"
+
+    def run_b(self, messages):
+        sys_msg = {"role": "system", "content": f"Predict next action from: {', '.join(ACTION_TYPES)}"}
+        out, i_t, o_t = self._gen(self.cheap_model, self.cheap_tok, [sys_msg] + messages)
+        return self._parse(out), i_t, o_t, "cheap"
+
+    def run_c(self, messages):
+        sys_msg = {"role": "system", "content": f"Predict next action from: {', '.join(ACTION_TYPES)}"}
+        proposal, i1, o1, _ = self._gen(self.cheap_model, self.cheap_tok, [sys_msg] + messages)
+        vp = messages + [{"role": "assistant", "content": proposal},
+                         {"role": "user", "content": "Is this action correct? Answer ONLY yes or no."}]
+        verdict, i2, o2, _ = self._gen(self.strong_model, self.strong_tok, vp, max_new=10)
+        if "yes" in verdict.lower():
+            return self._parse(proposal), i1 + i2, o1 + o2, "mixed"
+        out, i3, o3, _ = self._gen(self.strong_model, self.strong_tok, [sys_msg] + messages)
+        return self._parse(out), i1 + i2 + i3, o1 + o2 + o3, "mixed"
+
+    def run_d(self, messages):
+        if not self.verifier_name:
+            raise ValueError("Verifier model required for config D")
+        sys_msg = {"role": "system", "content": f"Predict next action from: {', '.join(ACTION_TYPES)}"}
+        proposal, i1, o1, _ = self._gen(self.cheap_model, self.cheap_tok, [sys_msg] + messages)
+        vp = messages + [{"role": "assistant", "content": proposal},
+                         {"role": "user", "content": "Rate this action: good or bad."}]
+        verdict, i2, o2, _ = self._gen(self.v_model, self.v_tok, vp, max_new=10)
+        if "good" in verdict.lower():
+            return self._parse(proposal), i1 + i2, o1 + o2, "cheap"
+        out, i3, o3, _ = self._gen(self.strong_model, self.strong_tok, [sys_msg] + messages)
+        return self._parse(out), i1 + i2 + i3, o1 + o2 + o3, "mixed"
+
+    def run_e(self, messages, n=3):
+        sys_msg = {"role": "system", "content": f"Predict next action from: {', '.join(ACTION_TYPES)}"}
+        proposals = []
+        total_i, total_o = 0, 0
+        for _ in range(n):
+            p, i_t, o_t, _ = self._gen(self.cheap_model, self.cheap_tok, [sys_msg] + messages, temp=0.7)
+            proposals.append(p)
+            total_i += i_t
+            total_o += o_t
+        best = proposals[0]
+        best_score = -1
+        for p in proposals:
+            rp = messages + [{"role": "assistant", "content": p},
+                             {"role": "user", "content": "Score this action 1-10."}]
+            s_text, i_t, o_t, _ = self._gen(self.strong_model, self.strong_tok, rp, max_new=5)
+            total_i += i_t
+            total_o += o_t
+            m = re.search(r'(\d+)', s_text)
+            if m:
+                sc = int(m.group(1))
+                if sc > best_score:
+                    best_score = sc
+                    best = p
+        return self._parse(best), total_i, total_o, "mixed"
+
+
+def evaluate(configs="ABCDE", limit=200):
+    print("\n=== Step 4: Evaluation ===")
+    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})
+    cost_log = []
+
+    for idx, ex in enumerate(ds):
+        msgs = ex["messages"]
+        gold = ex["action_type"]
+        for cfg in configs:
+            try:
+                if cfg == "A":
+                    pred, i_t, o_t, mtype = runner.run_a(msgs)
+                elif cfg == "B":
+                    pred, i_t, o_t, mtype = runner.run_b(msgs)
+                elif cfg == "C":
+                    pred, i_t, o_t, mtype = runner.run_c(msgs)
+                elif cfg == "D":
+                    pred, i_t, o_t, mtype = runner.run_d(msgs)
+                elif cfg == "E":
+                    pred, i_t, o_t, mtype = runner.run_e(msgs)
+                else:
+                    continue
+            except Exception as e:
+                print(f"Error {cfg} idx {idx}: {e}")
+                pred = "tool_call"
+                i_t, o_t, mtype = 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
+
+            cost = i_t * COST.get(f"{mtype}_in", 1.0) + o_t * COST.get(f"{mtype}_out", 1.0)
+            results[cfg]["cost"] += cost
+            cost_log.append({"config": cfg, "cost": cost})
+
+    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
+
+
+# ============================================================================
+# Step 5: Ablation Report + Cost-Quality Frontier
+# ============================================================================
+def generate_report(eval_results):
+    print("\n=== Step 5: Generating Report ===")
+    report = []
+    report.append("# Speculative Tool Actions — Ablation Report\n")
+    report.append("## Evaluation Configurations\n")
+    report.append("- **A**: Always strong model (Qwen2.5-7B)\n")
+    report.append("- **B**: Cheap model only (Qwen3-1.7B)\n")
+    report.append("- **C**: Cheap proposer + strong verifier\n")
+    report.append("- **D**: Cheap proposer + trained trace judge (Qwen3-4B reward model)\n")
+    report.append("- **E**: Multi-proposal reranking (3 cheap proposals + strong scoring)\n\n")
+
+    report.append("## Results\n\n")
+    report.append("| Config | Accuracy | Avg Cost | Unsafe-Action Rate |\n")
+    report.append("|--------|----------|----------|-------------------|\n")
+    for cfg in sorted(eval_results):
+        r = eval_results[cfg]
+        report.append(f"| {cfg} | {r['accuracy']:.3f} | {r['avg_cost']:.2f} | {r['unsafe_rate']:.3f} |\n")
+
+    report.append("\n## Cost-Quality Frontier\n\n")
+    # Pareto frontier: maximize accuracy per unit cost
+    points = []
+    for cfg, r in eval_results.items():
+        points.append((r["avg_cost"], r["accuracy"], cfg))
+    points.sort()
+    frontier = []
+    max_acc = -1
+    for cost, acc, cfg in points:
+        if acc > max_acc:
+            frontier.append((cost, acc, cfg))
+            max_acc = acc
+
+    report.append("Pareto-optimal configs (max accuracy for given cost):\n")
+    for cost, acc, cfg in frontier:
+        report.append(f"- **{cfg}**: cost={cost:.2f}, accuracy={acc:.3f}\n")
+
+    report.append("\n## Recommendations\n")
+    # Find best balance (highest accuracy / cost ratio)
+    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
+    report.append(f"- **Best accuracy/cost ratio**: Config {best_cfg} (ratio={best_ratio:.3f})\n")
+
+    # Find highest accuracy regardless of cost
+    best_acc_cfg = max(eval_results, key=lambda c: eval_results[c]["accuracy"])
+    report.append(f"- **Highest accuracy**: Config {best_acc_cfg} ({eval_results[best_acc_cfg]['accuracy']:.3f})\n")
+
+    # Find lowest cost with >90% of best accuracy
+    best_acc = eval_results[best_acc_cfg]["accuracy"]
+    threshold = best_acc * 0.9
+    cheap_candidates = {c: r for c, r in eval_results.items() if r["accuracy"] >= threshold}
+    if cheap_candidates:
+        cheapest = min(cheap_candidates, key=lambda c: cheap_candidates[c]["avg_cost"])
+        report.append(f"- **Cheapest config within 90% of best accuracy**: Config {cheapest} "
+                      f"(cost={cheap_candidates[cheapest]['avg_cost']:.2f}, accuracy={cheap_candidates[cheapest]['accuracy']:.3f})\n")
+
+    report_text = "".join(report)
+    with open("/tmp/ablation_report.md", "w") as f:
+        f.write(report_text)
+    print(report_text)
+    return report_text
+
+
+# ============================================================================
+# Main
+# ============================================================================
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--skip_build", action="store_true")
+    parser.add_argument("--skip_train_proposer", action="store_true")
+    parser.add_argument("--skip_train_verifier", action="store_true")
+    parser.add_argument("--skip_eval", action="store_true")
+    parser.add_argument("--eval_limit", type=int, default=200)
+    args = parser.parse_args()
+
+    if not args.skip_build:
+        build_datasets()
+    if not args.skip_train_proposer:
+        train_proposer()
+    if not args.skip_train_verifier:
+        train_verifier()
+    if not args.skip_eval:
+        results = evaluate(limit=args.eval_limit)
+        generate_report(results)
+    print("\nPipeline complete.")
+
+
+if __name__ == "__main__":
+    main()