205 lines
7.4 KiB
Python
205 lines
7.4 KiB
Python
#!/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())
|