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()