Organize Open-SWE-Traces data prep project
This commit is contained in:
Codex
204
scripts/repurposing/coarse_decompose.py
Normal file
204
scripts/repurposing/coarse_decompose.py
Normal 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())
|
||||
Reference in New Issue
Block a user