Add full end-to-end pipeline job script

pipeline_full.py

@@ -0,0 +1,436 @@
+#!/usr/bin/env python3
+"""
+Speculative Tool Actions — Full End-to-End Pipeline (single script)
+====================================================================
+Runs inside HF Jobs GPU instance. Steps:
+  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 configs A-E on held-out set
+  5. Generate cost-quality frontier + ablation report
+"""
+import os
+import sys
+import json
+import re
+import subprocess
+from collections import Counter, defaultdict
+from random import Random
+
+# Ensure deps
+subprocess.check_call([sys.executable, "-m", "pip", "install", "--quiet", "datasets", "transformers", "trl", "peft", "accelerate", "huggingface_hub", "trackio"])
+
+import torch
+from datasets import load_dataset, Dataset
+from transformers import AutoModelForCausalLM, AutoTokenizer, set_seed
+from trl import SFTTrainer, SFTConfig, RewardTrainer, RewardConfig
+from peft import LoraConfig
+
+set_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}
+
+# ========================================================================
+# Dataset building
+# ========================================================================
+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=3000, max_toolbench=2000):
+    print("\n=== STEP 1: Building Datasets ===")
+    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": [m.copy() for m in state], "completion": comp, "action_type": atype})
+                v_rows.append({"prompt": [m.copy() for m in state], "completion": comp, "label": bool(resolved), "action_type": atype})
+                e_rows.append({"messages": [m.copy() for m in state] + comp, "resolved": resolved, "action_type": atype})
+            state.append(msg)
+
+    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", {})
+        state = []
+        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.copy() for m in state], "completion": [msg.copy()], "action_type": atype})
+                v_rows.append({"prompt": [m.copy() for m in state], "completion": [msg.copy()], "label": True, "action_type": atype})
+                e_rows.append({"messages": [m.copy() for m in state] + [msg.copy()], "resolved": True, "action_type": atype})
+            state.append(msg)
+
+    print(f"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())
+
+    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 and 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("Pushed proposer 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": [m.copy() for m in 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": [m.copy() for m in 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("Pushed verifier dataset")
+
+    eval_ds = Dataset.from_list(e_rows).shuffle(seed=42).select(range(min(1000, len(e_rows))))
+    eval_ds.push_to_hub(f"{HUB_ORG}/speculative-actions-eval")
+    print("Pushed eval dataset")
+
+    return proposer_ds, verifier_ds, eval_ds
+
+
+# ========================================================================
+# Train proposer
+# ========================================================================
+def train_proposer():
+    print("\n=== STEP 2: Training Proposer ===")
+    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.")
+
+
+# ========================================================================
+# Train verifier
+# ========================================================================
+def train_verifier():
+    print("\n=== STEP 3: Training Verifier ===")
+    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
+        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)
+        return tokenizer.decode(out[0][inputs.shape[1]:], skip_special_tokens=True), 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=== 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})
+    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", runner.run_e)]:
+            try:
+                if cfg == "E":
+                    pred, i_t, o_t, mtype = func(msgs, n=3)
+                else:
+                    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) % 50 == 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=== STEP 5: Generating Report ===")
+    lines = ["# Speculative Tool Actions — Ablation Report\n"]
+    lines.append("## Configurations\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 + 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")
+    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():
+    proposer_ds, verifier_ds, eval_ds = build_datasets()
+    train_proposer()
+    train_verifier()
+    results = evaluate(limit=200)
+    generate_report(results)
+    print("\n=== Pipeline Complete ===")
+
+
+if __name__ == "__main__":
+    main()