Initial Open-SWE-Traces cleanup pipeline

2026-08-06 22:53:47 +08:00
commit 044bd03f0e
35 changed files with 3638 additions and 0 deletions
@@ -0,0 +1,61 @@
+#!/usr/bin/env python3
+"""Merge audit JSONL manifests into one unique source-ordered file."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+from pathlib import Path
+
+
+def main() -> None:
+    """Merge the last result for each sample and write it in source order."""
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("source", type=Path, help="Source sample JSONL defining output order")
+    parser.add_argument("output", type=Path)
+    parser.add_argument("manifests", type=Path, nargs="+")
+    parser.add_argument("--require-count", type=int)
+    args = parser.parse_args()
+
+    audits: dict[str, dict] = {}
+    for path in args.manifests:
+        if not path.exists():
+            continue
+        with path.open(encoding="utf-8") as handle:
+            for line in handle:
+                if not line.strip():
+                    continue
+                value = json.loads(line)
+                sample_id = value.get("sample_id")
+                if not isinstance(sample_id, str) or not sample_id:
+                    raise ValueError(f"Audit in {path} has no sample_id")
+                audits[sample_id] = value
+
+    ordered: list[dict] = []
+    with args.source.open(encoding="utf-8") as handle:
+        for line in handle:
+            if not line.strip():
+                continue
+            source_record = json.loads(line)
+            sample_id = source_record.get("trajectory_id") or source_record.get("sample_id")
+            if sample_id in audits:
+                ordered.append(audits.pop(sample_id))
+
+    if audits:
+        raise ValueError(f"Found {len(audits)} audit IDs absent from the source sample")
+    if args.require_count is not None and len(ordered) != args.require_count:
+        raise ValueError(f"Expected {args.require_count} merged records, found {len(ordered)}")
+
+    args.output.parent.mkdir(parents=True, exist_ok=True)
+    temporary = args.output.with_suffix(args.output.suffix + ".tmp")
+    with temporary.open("w", encoding="utf-8") as handle:
+        for value in ordered:
+            handle.write(json.dumps(value, ensure_ascii=False, default=str) + "\n")
+    os.replace(temporary, args.output)
+    print(json.dumps({"output": str(args.output), "records": len(ordered)}, indent=2))
+
+
+if __name__ == "__main__":
+    main()
@@ -0,0 +1,94 @@
+#!/usr/bin/env python3
+"""Create a deterministic uniform sample from selected Open-SWE outcome classes."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import random
+from pathlib import Path
+
+import pyarrow.parquet as pq
+
+
+def main() -> None:
+    """Sample complete rows uniformly from records matching the requested outcomes."""
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("dataset_dir", type=Path)
+    parser.add_argument("output", type=Path)
+    parser.add_argument("--count", type=int, default=50)
+    parser.add_argument("--seed", type=int, default=20260806)
+    parser.add_argument(
+        "--resolved",
+        type=int,
+        nargs="+",
+        default=[0, -1],
+        help="Outcome values eligible for sampling; defaults to failed and unknown trajectories.",
+    )
+    args = parser.parse_args()
+
+    files = sorted(args.dataset_dir.glob("data/**/*.parquet"))
+    if not files:
+        raise SystemExit(f"No parquet files found under {args.dataset_dir}")
+
+    allowed = set(args.resolved)
+    eligible: list[tuple[Path, int]] = []
+    for path in files:
+        # Reading only the outcome column keeps the full-dataset eligibility pass inexpensive.
+        outcomes = pq.read_table(path, columns=["resolved"])["resolved"].to_pylist()
+        eligible.extend((path, index) for index, value in enumerate(outcomes) if value in allowed)
+
+    if args.count > len(eligible):
+        raise SystemExit(f"Requested {args.count} rows from only {len(eligible)} eligible rows")
+
+    rng = random.Random(args.seed)
+    chosen = rng.sample(eligible, args.count)
+    selected: list[dict] = []
+    for path, file_row_index in chosen:
+        parquet_file = pq.ParquetFile(path)
+        row_group_start = 0
+        for row_group_index in range(parquet_file.num_row_groups):
+            row_group_rows = parquet_file.metadata.row_group(row_group_index).num_rows
+            if file_row_index < row_group_start + row_group_rows:
+                table = parquet_file.read_row_group(row_group_index)
+                row = table.slice(file_row_index - row_group_start, 1).to_pylist()[0]
+                row["_sample"] = {
+                    "sampling_method": "uniform without replacement over selected resolved values",
+                    "resolved_values": sorted(allowed),
+                    "seed": args.seed,
+                    "source_file": str(path.relative_to(args.dataset_dir)),
+                    "file_row_index": file_row_index,
+                    "row_group_index": row_group_index,
+                }
+                selected.append(row)
+                break
+            row_group_start += row_group_rows
+        else:
+            raise RuntimeError(f"Could not locate row {file_row_index} in {path}")
+
+    args.output.parent.mkdir(parents=True, exist_ok=True)
+    with args.output.open("w", encoding="utf-8") as handle:
+        for row in selected:
+            handle.write(json.dumps(row, ensure_ascii=False, default=str) + "\n")
+
+    outcome_counts = {str(value): 0 for value in sorted(allowed)}
+    for row in selected:
+        outcome_counts[str(row["resolved"])] += 1
+    print(
+        json.dumps(
+            {
+                "output": str(args.output),
+                "seed": args.seed,
+                "eligible_rows": len(eligible),
+                "sample_rows": len(selected),
+                "sample_outcome_counts": outcome_counts,
+            },
+            ensure_ascii=False,
+            indent=2,
+        )
+    )
+
+
+if __name__ == "__main__":
+    main()
@@ -0,0 +1,74 @@
+#!/usr/bin/env python3
+"""Split unfinished JSONL records into balanced temporary retry shards."""
+
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+
+
+def record_id(value: dict) -> str:
+    """Return the stable trajectory identifier used by audit manifests."""
+
+    sample_id = value.get("trajectory_id") or value.get("sample_id")
+    if not isinstance(sample_id, str) or not sample_id:
+        raise ValueError("JSONL record has no trajectory_id/sample_id")
+    return sample_id
+
+
+def main() -> None:
+    """Write unfinished input records round-robin across the requested shards."""
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("input", type=Path)
+    parser.add_argument("completed", type=Path)
+    parser.add_argument("output_dir", type=Path)
+    parser.add_argument("--shards", type=int, default=3)
+    args = parser.parse_args()
+    if args.shards < 1:
+        raise SystemExit("--shards must be positive")
+
+    completed_ids: set[str] = set()
+    with args.completed.open(encoding="utf-8") as handle:
+        for line in handle:
+            if line.strip():
+                completed_ids.add(record_id(json.loads(line)))
+
+    args.output_dir.mkdir(parents=True, exist_ok=True)
+    handles = [
+        (args.output_dir / f"part_{index}.input.jsonl").open("w", encoding="utf-8")
+        for index in range(args.shards)
+    ]
+    counts = [0] * args.shards
+    try:
+        unfinished = 0
+        with args.input.open(encoding="utf-8") as handle:
+            for line in handle:
+                if not line.strip():
+                    continue
+                value = json.loads(line)
+                if record_id(value) in completed_ids:
+                    continue
+                shard = unfinished % args.shards
+                handles[shard].write(line)
+                counts[shard] += 1
+                unfinished += 1
+    finally:
+        for handle in handles:
+            handle.close()
+
+    print(
+        json.dumps(
+            {
+                "completed": len(completed_ids),
+                "unfinished": sum(counts),
+                "shard_counts": counts,
+            },
+            indent=2,
+        )
+    )
+
+
+if __name__ == "__main__":
+    main()