| |
| """Convert public mini-swe-agent trajectories to Pi/Qwen structured tool chats.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import re |
| from pathlib import Path |
|
|
| from datasets import Dataset, load_dataset |
|
|
| SOURCE = "Kwai-Klear/SWE-smith-mini_swe_agent_plus-trajectories-66k" |
| BASH_BLOCK = re.compile(r"```(?:bash|sh)\s*\n(.*?)\n```", re.DOTALL | re.IGNORECASE) |
| FINAL_MARKER = "MINI_SWE_AGENT_FINAL_OUTPUT" |
|
|
| PI_SYSTEM_PROMPT = """You are an expert coding assistant operating inside pi, a coding agent harness. You help users by reading files, executing commands, editing code, and writing new files. |
| |
| Available tools: |
| - read: Read file contents |
| - bash: Execute bash commands |
| - edit: Make precise file edits with exact text replacement |
| - write: Create or overwrite files |
| |
| Guidelines: |
| - Inspect the repository and relevant instructions before changing code. |
| - Use tools to make the requested change; do not merely describe commands. |
| - Prefer focused searches and reads over dumping large files. |
| - Make minimal, targeted edits and preserve unrelated behavior. |
| - Run relevant focused tests after editing. |
| - When the task is complete, stop using tools and give a concise final summary. |
| - Be concise in your responses. |
| - Show file paths clearly when working with files.""" |
|
|
| PI_TOOLS = [ |
| { |
| "type": "function", |
| "function": { |
| "name": "read", |
| "description": "Read the contents of a file.", |
| "parameters": { |
| "type": "object", |
| "properties": { |
| "path": {"type": "string"}, |
| "offset": {"type": "number"}, |
| "limit": {"type": "number"}, |
| }, |
| "required": ["path"], |
| }, |
| }, |
| }, |
| { |
| "type": "function", |
| "function": { |
| "name": "bash", |
| "description": "Execute a bash command in the current working directory.", |
| "parameters": { |
| "type": "object", |
| "properties": { |
| "command": {"type": "string"}, |
| "timeout": {"type": "number"}, |
| }, |
| "required": ["command"], |
| }, |
| }, |
| }, |
| { |
| "type": "function", |
| "function": { |
| "name": "edit", |
| "description": "Edit one file using exact text replacement.", |
| "parameters": { |
| "type": "object", |
| "properties": { |
| "path": {"type": "string"}, |
| "edits": { |
| "type": "array", |
| "items": { |
| "type": "object", |
| "properties": { |
| "oldText": {"type": "string"}, |
| "newText": {"type": "string"}, |
| }, |
| "required": ["oldText", "newText"], |
| }, |
| }, |
| }, |
| "required": ["path", "edits"], |
| }, |
| }, |
| }, |
| { |
| "type": "function", |
| "function": { |
| "name": "write", |
| "description": "Create or overwrite a file.", |
| "parameters": { |
| "type": "object", |
| "properties": { |
| "path": {"type": "string"}, |
| "content": {"type": "string"}, |
| }, |
| "required": ["path", "content"], |
| }, |
| }, |
| }, |
| ] |
|
|
| TOOLS_JSON = json.dumps(PI_TOOLS, separators=(",", ":")) |
|
|
|
|
| def message( |
| role: str, |
| content: str, |
| *, |
| reasoning_content: str | None = None, |
| tool_calls: list[dict] | None = None, |
| tool_call_id: str | None = None, |
| name: str | None = None, |
| ) -> dict: |
| return { |
| "role": role, |
| "content": content, |
| "reasoning_content": reasoning_content, |
| "tool_calls": tool_calls, |
| "tool_call_id": tool_call_id, |
| "name": name, |
| } |
|
|
|
|
| def clean_reasoning(text: str) -> str: |
| text = text.strip() |
| text = re.sub(r"^THOUGHT:\s*", "", text, flags=re.IGNORECASE) |
| return text.strip() |
|
|
|
|
| def convert(example: dict) -> dict: |
| source_messages = example["messages"] |
| converted = [message("system", PI_SYSTEM_PROMPT)] |
| pending_call_id: str | None = None |
| call_index = 0 |
| initial_user_seen = False |
| malformed_actions = 0 |
|
|
| for source in source_messages: |
| role = source["role"] |
| content = source.get("content") or "" |
|
|
| if role == "system": |
| continue |
|
|
| if role == "user": |
| if pending_call_id is not None: |
| converted.append( |
| message( |
| "tool", |
| content, |
| tool_call_id=pending_call_id, |
| name="bash", |
| ) |
| ) |
| pending_call_id = None |
| else: |
| converted.append(message("user", content)) |
| initial_user_seen = True |
| continue |
|
|
| if role != "assistant": |
| continue |
|
|
| match = BASH_BLOCK.search(content) |
| is_final = FINAL_MARKER in content |
| if is_final: |
| converted.append( |
| message( |
| "assistant", |
| "Implemented the requested change and verified it with focused checks.", |
| ) |
| ) |
| pending_call_id = None |
| continue |
|
|
| if match is None: |
| malformed_actions += 1 |
| converted.append(message("assistant", content.strip())) |
| pending_call_id = None |
| continue |
|
|
| reasoning = clean_reasoning(content[: match.start()]) |
| command = match.group(1).strip() |
| call_index += 1 |
| pending_call_id = f"call_{call_index:04d}" |
| converted.append( |
| message( |
| "assistant", |
| "", |
| reasoning_content=reasoning, |
| tool_calls=[ |
| { |
| "id": pending_call_id, |
| "type": "function", |
| "function": { |
| "name": "bash", |
| "arguments": json.dumps( |
| {"command": command}, separators=(",", ":") |
| ), |
| }, |
| } |
| ], |
| ) |
| ) |
|
|
| if not initial_user_seen: |
| raise ValueError(f"trajectory {example['instance_id']} has no user request") |
|
|
| if converted[-1]["role"] != "assistant" or converted[-1]["tool_calls"]: |
| converted.append( |
| message( |
| "assistant", |
| "Implemented the requested change and verified it with focused checks.", |
| ) |
| ) |
|
|
| return { |
| "instance_id": example["instance_id"], |
| "messages": converted, |
| "tools": TOOLS_JSON, |
| "source": SOURCE, |
| "num_tool_calls": call_index, |
| "malformed_actions": malformed_actions, |
| } |
|
|
|
|
| def evaluation_instance_ids(path: Path) -> set[str]: |
| """Exact identifiers only, used solely as a decontamination blacklist.""" |
| if not path.exists(): |
| return set() |
| return {entry.name for entry in path.iterdir() if entry.is_dir()} |
|
|
|
|
| def write_shards(dataset: Dataset, output_dir: Path, shards: int) -> None: |
| output_dir.mkdir(parents=True, exist_ok=True) |
| for index in range(shards): |
| shard = dataset.shard(num_shards=shards, index=index, contiguous=True) |
| path = output_dir / f"train-{index:05d}-of-{shards:05d}.parquet" |
| shard.to_parquet(path) |
|
|
|
|
| def main() -> None: |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--output", type=Path, required=True) |
| parser.add_argument("--limit", type=int) |
| parser.add_argument("--num-proc", type=int, default=32) |
| parser.add_argument("--shards", type=int, default=16) |
| parser.add_argument( |
| "--eval-id-dir", |
| type=Path, |
| default=Path("/root/work/shared/tasksets/swe-bench-verified"), |
| ) |
| args = parser.parse_args() |
|
|
| split = "train" if args.limit is None else f"train[:{args.limit}]" |
| raw = load_dataset(SOURCE, split=split) |
| blacklist = evaluation_instance_ids(args.eval_id_dir) |
| before = len(raw) |
| raw = raw.filter( |
| lambda instance_id: instance_id not in blacklist, |
| input_columns=["instance_id"], |
| num_proc=args.num_proc, |
| desc="Removing exact evaluation identifiers", |
| ) |
| removed = before - len(raw) |
| converted = raw.map( |
| convert, |
| remove_columns=raw.column_names, |
| num_proc=args.num_proc, |
| desc="Converting trajectories to Pi structured calls", |
| ) |
|
|
| write_shards(converted, args.output, min(args.shards, len(converted))) |
| report = { |
| "source": SOURCE, |
| "source_revision": "750b2c11239fd5e32f97e6cfb9bf80fb9a9a2983", |
| "source_rows": before, |
| "rows_removed_as_exact_eval_ids": removed, |
| "output_rows": len(converted), |
| "tool_calls": sum(converted["num_tool_calls"]), |
| "malformed_actions": sum(converted["malformed_actions"]), |
| "transformation": "fenced bash -> Pi structured bash calls; terminal marker -> concise final response", |
| } |
| (args.output / "provenance.json").write_text(json.dumps(report, indent=2) + "\n") |
| print(json.dumps(report, indent=2)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
