Files
odysseus/scripts/assemble_sft_clean_corpus.py
T

75 lines
2.8 KiB
Python

#!/usr/bin/env python3
"""Assemble kept and validated repaired sessions into a clean SFT corpus."""
from __future__ import annotations
import argparse
import json
from collections import Counter, defaultdict
from pathlib import Path
from typing import Any
def load_jsonl(path: Path) -> list[dict[str, Any]]:
return [json.loads(line) for line in path.read_text(encoding="utf-8").splitlines() if line.strip()]
def main() -> None:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--trace", type=Path, required=True)
parser.add_argument("--verdicts", type=Path, required=True)
parser.add_argument("--repairs", type=Path, action="append", default=[])
parser.add_argument("--out-trace", type=Path, required=True)
parser.add_argument("--report", type=Path, required=True)
args = parser.parse_args()
source: dict[str, list[dict[str, Any]]] = defaultdict(list)
for row in load_jsonl(args.trace):
source[str(row.get("session_id") or "")].append(row)
verdicts = {str(row.get("session_id") or ""): row for row in load_jsonl(args.verdicts)}
repaired: dict[str, list[dict[str, Any]]] = defaultdict(list)
for path in args.repairs:
for row in load_jsonl(path):
repaired[str(row.get("session_id") or "")].append(row)
output: list[dict[str, Any]] = []
excluded: list[dict[str, Any]] = []
counts: Counter[str] = Counter()
for session_id in sorted(source):
verdict = verdicts.get(session_id)
decision = str((verdict or {}).get("verdict") or "missing")
if decision == "keep":
output.extend(source[session_id])
counts["kept"] += 1
elif decision == "repair" and repaired.get(session_id):
output.extend(repaired[session_id])
counts["repaired"] += 1
else:
counts["excluded"] += 1
excluded.append({
"session_id": session_id,
"verdict": decision,
"issues": (verdict or {}).get("issues") or [],
"repair_missing": decision == "repair" and session_id not in repaired,
})
args.out_trace.parent.mkdir(parents=True, exist_ok=True)
args.out_trace.write_text(
"\n".join(json.dumps(row, ensure_ascii=False) for row in output) + ("\n" if output else ""),
encoding="utf-8",
)
report = {
"source_sessions": len(source),
"output_sessions": counts["kept"] + counts["repaired"],
"output_turns": len(output),
"decisions": dict(counts),
"excluded": excluded,
}
args.report.parent.mkdir(parents=True, exist_ok=True)
args.report.write_text(json.dumps(report, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
print(json.dumps({key: value for key, value in report.items() if key != "excluded"}, indent=2))
if __name__ == "__main__":
main()