yi_lu/ti_coding_agent_data_prep
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Organize Open-SWE-Traces data prep project

Browse Source
This commit is contained in:
Codex
2026-06-24 22:34:04 +08:00
commit f06e573b04
27 changed files with 3264 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

259
scripts/repurposing/build_swift_training_probe_5k.py Normal file
View File

@@ -0,0 +1,259 @@
#!/usr/bin/env python3
from __future__ import annotations
import gzip
import json
import sys
from collections import Counter
from pathlib import Path
from typing import Any
import pyarrow as pa
import pyarrow.parquet as pq
REPO_ROOT = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(REPO_ROOT / "scripts" / "filtering"))
import audit_native_traces as audit # noqa: E402
TARGET_PER_CONFIG = 1250
NO_RESULT_TOOLS = {"finish", "submit", "think"}
CONFIG_META = {
"minimax_m25_openhands_trajectories": ("minimax_m25", "openhands", "thinking"),
"minimax_m25_sweagent_trajectories": ("minimax_m25", "sweagent", "thinking"),
"qwen35_openhands_trajectories": ("qwen35", "openhands", "non_thinking"),
"qwen35_sweagent_trajectories": ("qwen35", "sweagent", "non_thinking"),
}
class AuditArgs:
long_turn_threshold = 300
long_char_threshold = 900_000
repeat_window = 24
repeat_threshold = 10
def normalize_json(value: Any) -> Any:
return json.loads(json.dumps(value, ensure_ascii=False, default=str))
def parse_tool_args(raw: Any) -> Any:
if isinstance(raw, dict):
return raw
if isinstance(raw, str):
try:
return json.loads(raw)
except Exception:
return raw
return raw
def convert_tool_calls(tool_calls: list[dict[str, Any]] | None) -> list[dict[str, Any]]:
converted = []
for call in tool_calls or []:
function = call.get("function") or {}
converted.append(
{
"id": call.get("id"),
"type": call.get("type", "function"),
"function": {
"name": function.get("name") or call.get("name"),
"arguments": parse_tool_args(function.get("arguments", call.get("arguments"))),
},
}
)
return converted
def tool_calls_to_text(tool_calls: list[dict[str, Any]]) -> str:
if not tool_calls:
return ""
return "<tool_call>\n" + json.dumps(tool_calls, ensure_ascii=False, separators=(",", ":")) + "\n</tool_call>"
def tool_requires_result(tool_call: dict[str, Any]) -> bool:
name = (tool_call.get("function") or {}).get("name")
return name not in NO_RESULT_TOOLS
def convert_messages(trajectory: list[dict[str, Any]], thinking_mode: str, stats: Counter) -> list[dict[str, Any]]:
messages: list[dict[str, Any]] = []
pending_tool_calls: list[dict[str, Any]] = []
for source in trajectory:
role = source.get("role")
if role == "system":
content = source.get("content") or ""
if content:
messages.append({"role": "system", "content": content, "loss": False})
continue
if role == "user":
messages.append({"role": "user", "content": source.get("content") or "", "loss": False})
continue
if role == "assistant":
content = source.get("content") or ""
reasoning = source.get("reasoning_content") or ""
think_value = source.get("think")
if thinking_mode == "thinking":
if reasoning:
content = f"<think>\n{reasoning}\n</think>\n{content}"
stats["assistant_thinking_messages"] += 1 if reasoning else 0
else:
if reasoning:
stats["qwen_non_thinking_reasoning_nonempty"] += 1
if think_value not in (None, False):
stats[f"qwen_unexpected_think:{think_value}"] += 1
tool_calls = convert_tool_calls(source.get("tool_calls"))
tool_text = tool_calls_to_text(tool_calls)
if tool_text:
content = (content + "\n" + tool_text).strip()
msg = {"role": "assistant", "content": content, "loss": True}
if tool_calls:
msg["tool_calls"] = tool_calls
pending_tool_calls.extend([tool_call for tool_call in tool_calls if tool_requires_result(tool_call)])
messages.append(msg)
continue
if role == "tool":
tool_call = pending_tool_calls.pop(0) if pending_tool_calls else {}
messages.append(
{
"role": "tool",
"content": source.get("content") or "",
"tool_call_id": tool_call.get("id"),
"name": (tool_call.get("function") or {}).get("name"),
"loss": False,
}
)
continue
stats[f"unknown_role:{role}"] += 1
if pending_tool_calls:
stats["pending_tool_calls_after_conversion"] += len(pending_tool_calls)
return messages
def main() -> int:
input_root = Path("data/Open-SWE-Traces")
output_dir = Path("runs/training_probe_5k_swift")
output_dir.mkdir(parents=True, exist_ok=True)
audit_args = AuditArgs()
rows: list[dict[str, Any]] = []
stats = Counter()
for config in audit.CONFIGS:
model_family, scaffold, thinking_mode = CONFIG_META[config]
kept_for_config = 0
for file in sorted((input_root / "data" / config).glob("*.parquet")):
if kept_for_config >= TARGET_PER_CONFIG:
break
table = pq.read_table(file)
for row in table.to_pylist():
stats[f"seen:{config}"] += 1
issues, details = audit.audit_row(row, audit_args)
hard = audit.hard_filter_issues(issues)
if hard:
stats[f"filtered:{config}"] += 1
continue
messages = convert_messages(row.get("trajectory") or [], thinking_mode, stats)
item = {
"messages": messages,
"source_dataset": "nvidia/Open-SWE-Traces",
"source_config": config,
"model_family": model_family,
"scaffold": scaffold,
"thinking_mode": thinking_mode,
"instance_id": row.get("instance_id"),
"repo": row.get("repo"),
"language": row.get("language"),
"license": row.get("license"),
"trajectory_id": row.get("trajectory_id"),
"resolved": int(row.get("resolved") or 0),
"model_patch": row.get("model_patch") or "",
"metadata": normalize_json(row.get("metadata") or {}),
"audit_flags": issues,
"audit_details": dict(details),
}
rows.append(item)
kept_for_config += 1
stats[f"kept:{config}"] += 1
stats[f"resolved:{config}:{item['resolved']}"] += 1
if kept_for_config >= TARGET_PER_CONFIG:
break
jsonl = output_dir / "train.jsonl"
with jsonl.open("w", encoding="utf-8") as handle:
for row in rows:
handle.write(json.dumps(row, ensure_ascii=False, separators=(",", ":")) + "\n")
jsonl_gz = output_dir / "train.jsonl.gz"
with gzip.open(jsonl_gz, "wt", encoding="utf-8") as handle:
for row in rows:
handle.write(json.dumps(row, ensure_ascii=False, separators=(",", ":")) + "\n")
parquet_path = output_dir / "train.parquet"
parquet_rows = []
for row in rows:
parquet_rows.append(
{
**{k: v for k, v in row.items() if k not in {"messages", "metadata", "audit_flags", "audit_details"}},
"messages": json.dumps(row["messages"], ensure_ascii=False),
"metadata": json.dumps(row["metadata"], ensure_ascii=False),
"audit_flags": json.dumps(row["audit_flags"], ensure_ascii=False),
"audit_details": json.dumps(row["audit_details"], ensure_ascii=False),
}
)
pq.write_table(pa.Table.from_pylist(parquet_rows), parquet_path, compression="zstd")
metadata = {
"name": "open-swe-traces-swift-probe-5k",
"rows": len(rows),
"selection": f"{TARGET_PER_CONFIG} hard-filter-kept rows per source config",
"source_dataset": "nvidia/Open-SWE-Traces",
"format": "modelscope-swift messages JSONL",
"stats": dict(stats),
"thinking_policy": {
"minimax": "reasoning_content is wrapped as <think>...</think> in assistant content",
"qwen": "non-thinking export; reasoning_content is not emitted; unexpected nonempty reasoning is counted",
"tool_response_mask": "tool messages have loss=false",
},
"files": ["train.jsonl", "train.jsonl.gz", "train.parquet", "README.md", "metadata.json"],
}
(output_dir / "metadata.json").write_text(json.dumps(metadata, indent=2, ensure_ascii=False), encoding="utf-8")
(output_dir / "README.md").write_text(build_readme(metadata), encoding="utf-8")
print(json.dumps(metadata, indent=2, ensure_ascii=False))
return 0
def build_readme(metadata: dict[str, Any]) -> str:
return f"""---
license: other
task_categories:
- text-generation
language:
- code
pretty_name: Open-SWE-Traces Swift Probe 5K
size_categories:
- 1K<n<10K
---
# Open-SWE-Traces Swift Probe 5K
Balanced 5,000-row training probe subset from `nvidia/Open-SWE-Traces`, exported for ModelScope SWIFT-style SFT.
Selection:
- 1,250 hard-filter-kept rows from each source config.
- Original native scaffold semantics are preserved.
- MiniMax rows are exported as thinking examples by wrapping `reasoning_content` in `<think>...</think>`.
- Qwen rows are exported as non-thinking examples; `reasoning_content` is not emitted.
- Tool responses are included as `role: "tool"` messages with `loss: false`.
Stats:
```json
{json.dumps(metadata["stats"], indent=2, ensure_ascii=False)}
```
"""
if __name__ == "__main__":
raise SystemExit(main())

132
scripts/repurposing/build_swift_validation_500.py Normal file
View File

@@ -0,0 +1,132 @@
#!/usr/bin/env python3
from __future__ import annotations
import gzip
import json
import random
import sys
from collections import Counter
from pathlib import Path
from typing import Any
import pyarrow as pa
import pyarrow.parquet as pq
REPO_ROOT = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(REPO_ROOT / "scripts" / "filtering"))
import audit_native_traces as audit # noqa: E402
sys.path.insert(0, str(REPO_ROOT / "scripts" / "repurposing"))
import build_swift_training_probe_5k as train_builder # noqa: E402
TARGET_PER_CONFIG = 125
SEED = 20260624
def normalize_json(value: Any) -> Any:
return json.loads(json.dumps(value, ensure_ascii=False, default=str))
def load_train_ids(train_path: Path) -> set[str]:
ids = set()
with train_path.open("r", encoding="utf-8") as handle:
for line in handle:
row = json.loads(line)
ids.add(row["trajectory_id"])
return ids
def main() -> int:
input_root = Path("data/Open-SWE-Traces")
output_dir = Path("runs/training_probe_5k_swift")
train_ids = load_train_ids(output_dir / "train.jsonl")
audit_args = train_builder.AuditArgs()
rng = random.Random(SEED)
rows: list[dict[str, Any]] = []
stats = Counter()
for config in audit.CONFIGS:
candidates = []
model_family, scaffold, thinking_mode = train_builder.CONFIG_META[config]
for file in sorted((input_root / "data" / config).glob("*.parquet")):
table = pq.read_table(file)
for row in table.to_pylist():
stats[f"seen:{config}"] += 1
trajectory_id = row.get("trajectory_id")
if trajectory_id in train_ids:
stats[f"skip_train_overlap:{config}"] += 1
continue
issues, details = audit.audit_row(row, audit_args)
if audit.hard_filter_issues(issues):
stats[f"filtered:{config}"] += 1
continue
candidates.append((row, issues, details, model_family, scaffold, thinking_mode))
chosen = rng.sample(candidates, TARGET_PER_CONFIG)
for row, issues, details, model_family, scaffold, thinking_mode in chosen:
messages = train_builder.convert_messages(row.get("trajectory") or [], thinking_mode, stats)
item = {
"messages": messages,
"source_dataset": "nvidia/Open-SWE-Traces",
"source_config": config,
"model_family": model_family,
"scaffold": scaffold,
"thinking_mode": thinking_mode,
"instance_id": row.get("instance_id"),
"repo": row.get("repo"),
"language": row.get("language"),
"license": row.get("license"),
"trajectory_id": row.get("trajectory_id"),
"resolved": int(row.get("resolved") or 0),
"model_patch": row.get("model_patch") or "",
"metadata": normalize_json(row.get("metadata") or {}),
"audit_flags": issues,
"audit_details": dict(details),
}
rows.append(item)
stats[f"kept:{config}"] += 1
stats[f"resolved:{config}:{item['resolved']}"] += 1
jsonl = output_dir / "validation.jsonl"
with jsonl.open("w", encoding="utf-8") as handle:
for row in rows:
handle.write(json.dumps(row, ensure_ascii=False, separators=(",", ":")) + "\n")
with gzip.open(output_dir / "validation.jsonl.gz", "wt", encoding="utf-8") as handle:
for row in rows:
handle.write(json.dumps(row, ensure_ascii=False, separators=(",", ":")) + "\n")
parquet_rows = []
for row in rows:
parquet_rows.append(
{
**{k: v for k, v in row.items() if k not in {"messages", "metadata", "audit_flags", "audit_details"}},
"messages": json.dumps(row["messages"], ensure_ascii=False),
"metadata": json.dumps(row["metadata"], ensure_ascii=False),
"audit_flags": json.dumps(row["audit_flags"], ensure_ascii=False),
"audit_details": json.dumps(row["audit_details"], ensure_ascii=False),
}
)
pq.write_table(pa.Table.from_pylist(parquet_rows), output_dir / "validation.parquet", compression="zstd")
metadata_path = output_dir / "metadata.json"
metadata = json.loads(metadata_path.read_text(encoding="utf-8"))
metadata["validation"] = {
"rows": len(rows),
"selection": f"random {TARGET_PER_CONFIG} hard-filter-kept rows per source config, excluding train trajectory_id",
"seed": SEED,
"stats": dict(stats),
"files": ["validation.jsonl", "validation.jsonl.gz", "validation.parquet"],
}
metadata_path.write_text(json.dumps(metadata, indent=2, ensure_ascii=False), encoding="utf-8")
readme = output_dir / "README.md"
text = readme.read_text(encoding="utf-8")
text += f"\n\n## Validation Split\n\nRandom balanced validation split, seed `{SEED}`, 125 rows per config, excluding train trajectory ids.\n\n```json\n{json.dumps(metadata['validation']['stats'], indent=2, ensure_ascii=False)}\n```\n"
readme.write_text(text, encoding="utf-8")
print(json.dumps(metadata["validation"], indent=2, ensure_ascii=False))
return 0
if __name__ == "__main__":
raise SystemExit(main())

171
scripts/repurposing/build_training_probe_5k.py Normal file
View File

@@ -0,0 +1,171 @@
#!/usr/bin/env python3
from __future__ import annotations
import gzip
import json
import sys
from collections import Counter
from pathlib import Path
from typing import Any
import pyarrow as pa
import pyarrow.parquet as pq
REPO_ROOT = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(REPO_ROOT / "scripts" / "filtering"))
import audit_native_traces as audit # noqa: E402
class AuditArgs:
long_turn_threshold = 300
long_char_threshold = 900_000
repeat_window = 24
repeat_threshold = 10
CONFIG_META = {
"minimax_m25_openhands_trajectories": ("minimax_m25", "openhands", "thinking"),
"minimax_m25_sweagent_trajectories": ("minimax_m25", "sweagent", "thinking"),
"qwen35_openhands_trajectories": ("qwen35", "openhands", "non_thinking"),
"qwen35_sweagent_trajectories": ("qwen35", "sweagent", "non_thinking"),
}
def normalize_for_json(value: Any) -> Any:
return json.loads(json.dumps(value, ensure_ascii=False, default=str))
def main() -> int:
input_root = Path("data/Open-SWE-Traces")
output_dir = Path("runs/training_probe_5k")
target = 5000
output_dir.mkdir(parents=True, exist_ok=True)
rows: list[dict[str, Any]] = []
stats = Counter()
audit_args = AuditArgs()
for config in audit.CONFIGS:
model_family, scaffold, thinking_mode = CONFIG_META[config]
for file in sorted((input_root / "data" / config).glob("*.parquet")):
table = pq.read_table(file)
for row in table.to_pylist():
issues, details = audit.audit_row(row, audit_args)
hard = audit.hard_filter_issues(issues)
stats["seen"] += 1
if hard:
stats["filtered"] += 1
continue
item = {
"probe_index": len(rows),
"source_dataset": "nvidia/Open-SWE-Traces",
"source_config": config,
"model_family": model_family,
"scaffold": scaffold,
"thinking_mode": thinking_mode,
"instance_id": row.get("instance_id"),
"repo": row.get("repo"),
"license": row.get("license"),
"language": row.get("language"),
"trajectory_id": row.get("trajectory_id"),
"trajectory": normalize_for_json(row.get("trajectory") or []),
"model_patch": row.get("model_patch") or "",
"resolved": int(row.get("resolved") or 0),
"metadata": normalize_for_json(row.get("metadata") or {}),
"audit_flags": issues,
"hard_filter_issues": hard,
"audit_details": dict(details),
}
rows.append(item)
stats["kept"] += 1
stats[f"config:{config}"] += 1
stats[f"resolved:{item['resolved']}"] += 1
if len(rows) >= target:
break
if len(rows) >= target:
break
if len(rows) >= target:
break
jsonl_gz = output_dir / "train.jsonl.gz"
with gzip.open(jsonl_gz, "wt", encoding="utf-8") as handle:
for row in rows:
handle.write(json.dumps(row, ensure_ascii=False, separators=(",", ":")) + "\n")
parquet_path = output_dir / "train.parquet"
table = pa.Table.from_pylist(rows)
pq.write_table(table, parquet_path, compression="zstd")
metadata = {
"name": "open-swe-traces-native-probe-5k",
"source_dataset": "nvidia/Open-SWE-Traces",
"selection": "first 5000 hard-filter-kept traces in config/shard order",
"rows": len(rows),
"stats": dict(stats),
"hard_filter_policy": {
"script": "scripts/filtering/audit_native_traces.py",
"long_turn_threshold": audit_args.long_turn_threshold,
"long_char_threshold": audit_args.long_char_threshold,
"repeat_window": audit_args.repeat_window,
"repeat_threshold": audit_args.repeat_threshold,
},
"files": {
"jsonl_gz": str(jsonl_gz),
"parquet": str(parquet_path),
},
}
(output_dir / "metadata.json").write_text(json.dumps(metadata, indent=2, ensure_ascii=False), encoding="utf-8")
(output_dir / "README.md").write_text(build_readme(metadata), encoding="utf-8")
print(json.dumps(metadata, indent=2, ensure_ascii=False))
return 0
def build_readme(metadata: dict[str, Any]) -> str:
return f"""---
license: other
task_categories:
- text-generation
language:
- code
pretty_name: Open-SWE-Traces Native Probe 5K
size_categories:
- 1K<n<10K
---
# Open-SWE-Traces Native Probe 5K
This is a 5,000-row training probe subset derived from `nvidia/Open-SWE-Traces`.
Selection policy:
- Keep the original native scaffold trajectory format.
- Iterate configs and parquet shards in sorted order.
- Apply the local hard-filter audit used in `audit_openswe_native_traces.py`.
- Take the first 5,000 rows that pass the hard filter.
Included fields:
- `source_config`
- `model_family`
- `scaffold`
- `thinking_mode`
- `instance_id`
- `repo`
- `trajectory`
- `model_patch`
- `resolved`
- `metadata`
- `audit_flags`
- `audit_details`
Rows: {metadata["rows"]}
Stats:
```json
{json.dumps(metadata["stats"], indent=2, ensure_ascii=False)}
```
"""
if __name__ == "__main__":
raise SystemExit(main())

204
scripts/repurposing/coarse_decompose.py Normal file
View File

@@ -0,0 +1,204 @@
#!/usr/bin/env python3
from __future__ import annotations
import json
from collections import Counter
from pathlib import Path
from typing import Any
IN_FILE = Path("runs/subproblem_decomposition/random20_decomposed.json")
OUT_FILE = Path("runs/subproblem_decomposition/random20_coarse_decomposed.json")
REPORT_FILE = Path("runs/subproblem_decomposition/coarse_report.json")
COARSE_MAP = {
"understand_task": "understand",
"explore_repo": "locate",
"locate_relevant_code": "locate",
"inspect_code": "diagnose",
"reproduce_or_probe": "diagnose",
"edit_solution": "fix",
"verify_solution": "verify",
"cleanup": "finalize",
"review_diff": "finalize",
"submit": "finalize",
"other": "diagnose",
}
COARSE_GOALS = {
"understand": "Understand the issue, repository context, constraints, and success criteria.",
"locate": "Find the files, modules, symbols, or tests likely responsible for the issue.",
"diagnose": "Inspect behavior and evidence to determine the concrete root cause.",
"fix": "Edit the implementation or tests to address the root cause.",
"verify": "Run targeted checks to confirm the fix and catch regressions.",
"finalize": "Clean temporary artifacts, inspect the final diff, and submit the solution.",
}
def dedupe(items: list[Any], limit: int | None = None) -> list[Any]:
seen = set()
out = []
for item in items:
key = json.dumps(item, ensure_ascii=False, sort_keys=True) if isinstance(item, (dict, list)) else str(item)
if key in seen:
continue
seen.add(key)
out.append(item)
if limit and len(out) >= limit:
break
return out
def merge_segments(segments: list[dict[str, Any]]) -> list[dict[str, Any]]:
"""Aggregate fine segments into a small fixed set of task stages.
Fine-grained traces often alternate between locate/diagnose/fix many times.
For SFT subproblem construction we want coarse, purposeful stages rather
than exact chronological micro-steps, so this groups by task category and
preserves the first-to-last turn span for each category.
"""
by_category: dict[str, dict[str, Any]] = {}
def start(seg: dict[str, Any], name: str) -> dict[str, Any]:
name = COARSE_MAP.get(seg["phase"], "diagnose")
return {
"category": name,
"goal": COARSE_GOALS[name],
"turn_range": list(seg["turn_range"]),
"source_phases": [seg["phase"]],
"commands": list(seg.get("commands", [])),
"tool_names": list(seg.get("tool_names", [])),
"files": list(seg.get("files", [])),
"assistant_intents": list(seg.get("assistant_intents", [])),
"observations": list(seg.get("observations", [])),
}
for seg in segments:
name = COARSE_MAP.get(seg["phase"], "diagnose")
current = by_category.get(name)
if current is None:
by_category[name] = start(seg, name)
continue
current["turn_range"][0] = min(current["turn_range"][0], seg["turn_range"][0])
current["turn_range"][1] = max(current["turn_range"][1], seg["turn_range"][1])
current["source_phases"].append(seg["phase"])
current["commands"].extend(seg.get("commands", []))
current["tool_names"].extend(seg.get("tool_names", []))
current["files"].extend(seg.get("files", []))
current["assistant_intents"].extend(seg.get("assistant_intents", []))
current["observations"].extend(seg.get("observations", []))
ordered = []
for name in ["understand", "locate", "diagnose", "fix", "verify", "finalize"]:
if name in by_category:
ordered.append(finish(by_category[name]))
return ordered
def finish(seg: dict[str, Any]) -> dict[str, Any]:
seg["source_phases"] = dedupe(seg["source_phases"])
seg["commands"] = dedupe(seg["commands"], 12)
seg["tool_names"] = dedupe(seg["tool_names"], 12)
seg["files"] = dedupe(seg["files"], 12)
seg["assistant_intents"] = dedupe(seg["assistant_intents"], 6)
seg["observations"] = dedupe(seg["observations"], 6)
return seg
def build_training_views(item: dict[str, Any], coarse_segments: list[dict[str, Any]]) -> dict[str, Any]:
qa = []
question_trajectory = []
for seg in coarse_segments:
qa.append(
{
"question": f"In repository {item['repo']}, how should an agent {seg['goal'][0].lower() + seg['goal'][1:]}",
"answer": {
"category": seg["category"],
"files": seg["files"],
"commands": seg["commands"],
"evidence": seg["observations"],
"next_action_hint": next_action_hint(seg["category"]),
},
}
)
question_trajectory.append(
{
"question": f"For task {item['instance_id']} in {item['repo']}, perform the {seg['category']} stage.",
"category": seg["category"],
"turn_range": seg["turn_range"],
"source_phases": seg["source_phases"],
"trajectory_outline": {
"intents": seg["assistant_intents"],
"tools": seg["tool_names"],
"commands": seg["commands"],
"observations": seg["observations"],
},
}
)
return {"qa": qa, "question_trajectory": question_trajectory}
def next_action_hint(category: str) -> str:
return {
"understand": "Start repository exploration.",
"locate": "Inspect the most relevant code paths.",
"diagnose": "Decide the minimal edit needed to fix the root cause.",
"fix": "Run targeted verification.",
"verify": "Review diff and clean temporary files.",
"finalize": "Stop or submit the final solution.",
}[category]
def main() -> int:
items = json.loads(IN_FILE.read_text())
output = []
category_counts = Counter()
before = 0
after = 0
for item in items:
coarse_segments = merge_segments(item["segments"])
before += len(item["segments"])
after += len(coarse_segments)
category_counts.update(seg["category"] for seg in coarse_segments)
converted = {
**{k: item[k] for k in [
"config",
"scaffold",
"model_family",
"thinking_mode",
"instance_id",
"repo",
"language",
"license",
"trajectory_id",
"resolved",
"trajectory_len",
"task",
"model_patch_head",
]},
"coarse_segments": coarse_segments,
"training_views": build_training_views(item, coarse_segments),
}
output.append(converted)
report = {
"items": len(items),
"fine_segments": before,
"coarse_segments": after,
"avg_fine_segments_per_trace": before / len(items),
"avg_coarse_segments_per_trace": after / len(items),
"category_counts": dict(category_counts),
"categories": COARSE_GOALS,
}
OUT_FILE.write_text(json.dumps(output, indent=2, ensure_ascii=False))
REPORT_FILE.write_text(json.dumps(report, indent=2, ensure_ascii=False))
print(json.dumps(report, indent=2, ensure_ascii=False))
return 0
if __name__ == "__main__":
raise SystemExit(main())

419
scripts/repurposing/convert_openswe_to_pi_mono.py Normal file
View File

@@ -0,0 +1,419 @@
#!/usr/bin/env python3
from __future__ import annotations
import argparse
import gzip
import json
import re
from collections import Counter
from pathlib import Path
from typing import Any
import pyarrow.parquet as pq
CONFIGS = {
"minimax_m25_openhands_trajectories": {
"model_family": "minimax_m25",
"scaffold": "openhands",
"thinking_mode": "thinking",
},
"minimax_m25_sweagent_trajectories": {
"model_family": "minimax_m25",
"scaffold": "sweagent",
"thinking_mode": "thinking",
},
"qwen35_openhands_trajectories": {
"model_family": "qwen35",
"scaffold": "openhands",
"thinking_mode": "non_thinking",
},
"qwen35_sweagent_trajectories": {
"model_family": "qwen35",
"scaffold": "sweagent",
"thinking_mode": "non_thinking",
},
}
PI_TOOL_NAMES = {"bash", "read", "edit", "write", "grep", "find", "ls"}
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description="Convert Open-SWE-Traces to pi-mono message JSONL.")
parser.add_argument("--input-root", type=Path, default=Path("data/Open-SWE-Traces"))
parser.add_argument("--output-root", type=Path, default=Path("runs/pi_mono_converted"))
parser.add_argument("--limit", type=int, default=0, help="Optional total row limit for smoke tests.")
parser.add_argument("--sample", type=int, default=20, help="Number of converted examples to save as sample JSON.")
parser.add_argument("--batch-size", type=int, default=256)
return parser.parse_args()
def repo_workspace_name(repo: str) -> str:
return repo.replace("/", "__") + "__1.0"
def root_patterns(row: dict[str, Any]) -> list[str]:
repo = row.get("repo") or ""
patterns = ["/testbed"]
if repo:
patterns.append("/workspace/" + repo_workspace_name(repo))
uploaded = extract_uploaded_path(row.get("trajectory") or [])
if uploaded:
patterns.append(uploaded.rstrip("/"))
return dedupe([p for p in patterns if p])
def extract_uploaded_path(trajectory: list[dict[str, Any]]) -> str | None:
for msg in trajectory:
if msg.get("role") != "user":
continue
content = msg.get("content") or ""
match = re.search(r"<uploaded_files>\s*(.*?)\s*</uploaded_files>", content, re.S)
if match:
return match.group(1).strip().splitlines()[0].rstrip("/")
return None
def dedupe(items: list[Any]) -> list[Any]:
seen = set()
out = []
for item in items:
key = json.dumps(item, sort_keys=True, ensure_ascii=False) if isinstance(item, (dict, list)) else str(item)
if key in seen:
continue
seen.add(key)
out.append(item)
return out
def normalize_text_paths(text: str, roots: list[str]) -> str:
if not text:
return text
out = text
for root in sorted(roots, key=len, reverse=True):
escaped = re.escape(root.rstrip("/"))
out = re.sub(escaped + r"(?=/|\b|$)", ".", out)
out = re.sub(r"<uploaded_files>\s*\.\s*</uploaded_files>", "<uploaded_files>\n.\n</uploaded_files>", out)
return out
def normalize_path(path: Any, roots: list[str]) -> str:
if not isinstance(path, str) or not path:
return "."
cleaned = normalize_text_paths(path.strip(), roots)
cleaned = cleaned.rstrip("/")
if cleaned in {"", "."}:
return "."
if cleaned.startswith("./"):
return cleaned[2:] or "."
return cleaned
def parse_tool_args(raw_args: Any, warnings: list[str]) -> dict[str, Any]:
if isinstance(raw_args, dict):
return raw_args
if not isinstance(raw_args, str):
warnings.append("tool_arguments_not_string_or_dict")
return {}
try:
parsed = json.loads(raw_args)
if isinstance(parsed, dict):
return parsed
warnings.append("tool_arguments_json_not_object")
return {}
except Exception:
warnings.append("tool_arguments_json_parse_failed")
return {}
def next_tool_observation(trajectory: list[dict[str, Any]], index: int) -> str:
for msg in trajectory[index + 1 :]:
if msg.get("role") == "tool":
return msg.get("content") or ""
if msg.get("role") == "assistant":
return ""
return ""
def looks_like_file_view(path: str, args: dict[str, Any], observation: str) -> bool:
if args.get("view_range"):
return True
if "cat -n" in observation or "result of running" in observation:
return True
if "files and directories" in observation:
return False
name = path.rsplit("/", 1)[-1]
if "." in name:
return True
known_files = {"Makefile", "Dockerfile", "LICENSE", "README", "Rakefile", "Gemfile", "Cargo.toml", "go.mod"}
return name in known_files
def line_range_to_offset_limit(view_range: Any) -> dict[str, int]:
if not isinstance(view_range, list) or len(view_range) != 2:
return {}
start, end = view_range
if not isinstance(start, int) or not isinstance(end, int):
return {}
offset = max(start - 1, 0)
if end < start:
return {"offset": offset}
return {"offset": offset, "limit": end - start + 1}
def convert_tool_call(
call: dict[str, Any],
roots: list[str],
observation: str,
warnings: list[str],
) -> tuple[dict[str, Any] | None, str | None, bool]:
function = call.get("function") or {}
source_name = function.get("name") or call.get("name") or ""
source_args = parse_tool_args(function.get("arguments", call.get("arguments")), warnings)
call_id = str(call.get("id") or f"call_{abs(hash(json.dumps(call, sort_keys=True, default=str))) % 10**12}")
if source_name in {"execute_bash", "bash"}:
command = normalize_text_paths(str(source_args.get("command") or ""), roots)
return {"type": "toolCall", "id": call_id, "name": "bash", "arguments": {"command": command}}, None, True
if source_name == "str_replace_editor":
command = source_args.get("command")
path = normalize_path(source_args.get("path"), roots)
if command == "view":
if looks_like_file_view(path, source_args, observation):
args = {"path": path}
args.update(line_range_to_offset_limit(source_args.get("view_range")))
return {"type": "toolCall", "id": call_id, "name": "read", "arguments": args}, None, True
return {"type": "toolCall", "id": call_id, "name": "ls", "arguments": {"path": path}}, None, True
if command == "create":
content = normalize_text_paths(str(source_args.get("file_text") or ""), roots)
return {"type": "toolCall", "id": call_id, "name": "write", "arguments": {"path": path, "content": content}}, None, True
if command == "str_replace":
old_text = normalize_text_paths(str(source_args.get("old_str") or ""), roots)
new_text = normalize_text_paths(str(source_args.get("new_str") or ""), roots)
return {
"type": "toolCall",
"id": call_id,
"name": "edit",
"arguments": {"path": path, "edits": [{"oldText": old_text, "newText": new_text}]},
}, None, True
if command == "insert":
warnings.append("str_replace_editor_insert_mapped_to_bash")
insert_line = source_args.get("insert_line")
new_text = source_args.get("new_str") or source_args.get("insert_str") or ""
py = (
"python - <<'PY'\n"
"from pathlib import Path\n"
f"path = Path({path!r})\n"
f"line = {int(insert_line) if isinstance(insert_line, int) else 0}\n"
f"text = {new_text!r}\n"
"lines = path.read_text().splitlines(True)\n"
"idx = max(0, min(line, len(lines)))\n"
"lines[idx:idx] = [text if text.endswith('\\n') else text + '\\n']\n"
"path.write_text(''.join(lines))\n"
"PY"
)
return {"type": "toolCall", "id": call_id, "name": "bash", "arguments": {"command": py}}, None, True
warnings.append(f"unsupported_str_replace_editor_command:{command}")
return None, None, True
if source_name in {"finish", "submit"}:
text = source_args.get("answer") or source_args.get("summary") or source_args.get("message") or ""
return None, normalize_text_paths(str(text), roots), True
if source_name == "think":
thought = source_args.get("thought") or source_args.get("content") or ""
return None, normalize_text_paths(str(thought), roots), True
warnings.append(f"unsupported_tool:{source_name}")
return None, None, True
def clean_tool_result(content: str, roots: list[str]) -> str:
text = normalize_text_paths(content or "", roots)
text = re.sub(r"^OBSERVATION:\s*\n?", "", text)
return text
def convert_row(row: dict[str, Any], config_name: str) -> dict[str, Any]:
config = CONFIGS[config_name]
warnings: list[str] = []
roots = root_patterns(row)
trajectory = row.get("trajectory") or []
pending: list[dict[str, Any]] = []
skipped_tool_results = 0
messages: list[dict[str, Any]] = []
system_messages: list[str] = []
for index, msg in enumerate(trajectory):
role = msg.get("role")
if role == "system":
content = normalize_text_paths(msg.get("content") or "", roots)
if content:
system_messages.append(content)
continue
if role == "user":
messages.append({"role": "user", "content": normalize_text_paths(msg.get("content") or "", roots)})
continue
if role == "assistant":
blocks: list[dict[str, Any]] = []
reasoning = msg.get("reasoning_content") or ""
if config["thinking_mode"] == "thinking" and reasoning.strip():
blocks.append({"type": "thinking", "thinking": normalize_text_paths(reasoning, roots)})
visible = msg.get("content") or ""
if visible.strip():
blocks.append({"type": "text", "text": normalize_text_paths(visible, roots)})
final_texts: list[str] = []
observation = next_tool_observation(trajectory, index)
for call in msg.get("tool_calls") or []:
converted, final_text, consumes_result = convert_tool_call(call, roots, observation, warnings)
if converted is not None:
blocks.append(converted)
pending.append(converted)
elif consumes_result:
skipped_tool_results += 1
if final_text:
final_texts.append(final_text)
for text in final_texts:
if text.strip():
blocks.append({"type": "text", "text": text})
if blocks:
stop_reason = "toolUse" if any(block.get("type") == "toolCall" for block in blocks) else "stop"
messages.append({"role": "assistant", "content": blocks, "stopReason": stop_reason})
continue
if role == "tool":
if not pending:
if skipped_tool_results > 0:
skipped_tool_results -= 1
continue
warnings.append("tool_result_without_pending_call")
continue
call = pending.pop(0)
messages.append(
{
"role": "toolResult",
"toolCallId": call["id"],
"toolName": call["name"],
"content": [{"type": "text", "text": clean_tool_result(msg.get("content") or "", roots)}],
"isError": infer_tool_error(msg.get("content") or ""),
}
)
continue
warnings.append(f"unsupported_role:{role}")
if pending:
warnings.append(f"pending_tool_calls_unmatched:{len(pending)}")
return {
"instance_id": row.get("instance_id"),
"repo": row.get("repo"),
"license": row.get("license"),
"language": row.get("language"),
"trajectory_id": row.get("trajectory_id"),
"config": config_name,
"scaffold": config["scaffold"],
"model_family": config["model_family"],
"thinking_mode": config["thinking_mode"],
"resolved": int(row.get("resolved") or 0),
"metadata": row.get("metadata") or {},
"repo_roots": roots,
"system_messages": system_messages,
"pi_mono_messages": messages,
"model_patch": normalize_text_paths(row.get("model_patch") or "", roots),
"conversion_warnings": dedupe(warnings),
}
def infer_tool_error(content: str) -> bool:
text = content or ""
error_markers = [
"Traceback (most recent call last)",
"Command failed",
"exit code 1",
"Invalid `view_range`",
"No such file or directory",
]
return any(marker in text for marker in error_markers)
def write_jsonl_gz(path: Path, rows: list[dict[str, Any]]) -> None:
path.parent.mkdir(parents=True, exist_ok=True)
with gzip.open(path, "at", encoding="utf-8") as handle:
for row in rows:
handle.write(json.dumps(row, ensure_ascii=False, separators=(",", ":")) + "\n")
def convert_file(file: Path, out_file: Path, config_name: str, batch_size: int, remaining: int | None) -> dict[str, Any]:
stats = Counter()
samples: list[dict[str, Any]] = []
parquet = pq.ParquetFile(file)
if out_file.exists():
out_file.unlink()
for batch in parquet.iter_batches(batch_size=batch_size):
rows = batch.to_pylist()
if remaining is not None:
rows = rows[:remaining]
converted = [convert_row(row, config_name) for row in rows]
write_jsonl_gz(out_file, converted)
for item in converted:
stats["rows"] += 1
stats["messages"] += len(item["pi_mono_messages"])
stats["warnings"] += len(item["conversion_warnings"])
stats[f"resolved_{item['resolved']}"] += 1
for warning in item["conversion_warnings"]:
stats[f"warning:{warning}"] += 1
if len(samples) < 3:
samples.append(item)
if remaining is not None:
remaining -= len(rows)
if remaining <= 0:
break
return {"stats": dict(stats), "samples": samples}
def main() -> int:
args = parse_args()
data_root = args.input_root / "data"
args.output_root.mkdir(parents=True, exist_ok=True)
all_stats: dict[str, Any] = {
"input_root": str(args.input_root),
"output_root": str(args.output_root),
"configs": {},
}
samples: list[dict[str, Any]] = []
remaining = args.limit or None
for config_name in CONFIGS:
config_dir = data_root / config_name
out_dir = args.output_root / config_name
out_dir.mkdir(parents=True, exist_ok=True)
config_stats = Counter()
for file in sorted(config_dir.glob("*.parquet")):
if remaining is not None and remaining <= 0:
break
out_file = out_dir / (file.stem + ".pi_mono.jsonl.gz")
result = convert_file(file, out_file, config_name, args.batch_size, remaining)
file_stats = Counter(result["stats"])
config_stats.update(file_stats)
samples.extend(result["samples"])
if remaining is not None:
remaining -= file_stats["rows"]
print(json.dumps({"file": str(file), "out": str(out_file), "stats": dict(file_stats)}, ensure_ascii=False))
all_stats["configs"][config_name] = dict(config_stats)
report_path = args.output_root / "conversion_report.json"
report_path.write_text(json.dumps(all_stats, indent=2, ensure_ascii=False), encoding="utf-8")
sample_path = args.output_root / "sample_converted.json"
sample_path.write_text(json.dumps(samples[: args.sample], indent=2, ensure_ascii=False), encoding="utf-8")
print(json.dumps({"report": str(report_path), "sample": str(sample_path)}, ensure_ascii=False))
return 0
if __name__ == "__main__":
raise SystemExit(main())