#!/usr/bin/env python3 import argparse import hashlib import json from pathlib import Path from build_case_pack_simulation_inputs_v1 import ( COMMON_INSTRUCTIONS, OUTPUT_SCHEMA, dump, iter_jsonl, load_case_pack, load_extensions, short_persona, ) PERSONA_FILE = Path("data/processed/personas_full_v1_0.jsonl") ORIGINAL_MANIFEST = Path("data/experimental/arca_nagano_reports/case_packs_v1/case_pack_manifest.jsonl") REPLACEMENT_MANIFEST = Path("data/experimental/arca_nagano_reports/replacements_20260622/case_pack_manifest.jsonl") INVALID_OLD_REPORT_IDS = { "05_shhxgwd_case_v1", "21_s4v7pyc_case_v1", "28_sak6a95_case_v1", "29_sfzmfhy_case_v1", } def read_pairs(path: Path) -> set[tuple[str, str]]: pairs: set[tuple[str, str]] = set() if not path or not path.exists(): return pairs with path.open(encoding="utf-8") as handle: for line in handle: if not line.strip(): continue parts = line.rstrip("\n").split("\t") if len(parts) >= 2: pairs.add((parts[0], parts[1])) return pairs def write_combined_manifest(path: Path) -> list[dict]: rows = [] for row in iter_jsonl(ORIGINAL_MANIFEST): if row["case_pack_id"] not in INVALID_OLD_REPORT_IDS: rows.append(row) rows.extend(iter_jsonl(REPLACEMENT_MANIFEST)) rows.sort(key=lambda row: (int(row.get("rank", 999)), row["case_pack_id"])) path.parent.mkdir(parents=True, exist_ok=True) with path.open("w", encoding="utf-8") as out: for row in rows: out.write(dump(row) + "\n") return rows def main(): parser = argparse.ArgumentParser() parser.add_argument("--run-id", required=True) parser.add_argument("--outdir", required=True) parser.add_argument("--workers", required=True, help="Comma-separated shard names") parser.add_argument( "--combined-manifest", default="data/experimental/arca_nagano_reports/case_packs_v1/case_pack_manifest_eval30_replaced_20260704.jsonl", ) parser.add_argument("--completed-pairs", default="", help="TSV: report_idpersona_id to skip") parser.add_argument("--target-pairs", default="", help="TSV: only these report_id/persona_id pairs") parser.add_argument("--model", default="deepseek-v4-flash") args = parser.parse_args() outdir = Path(args.outdir) outdir.mkdir(parents=True, exist_ok=True) workers = [w.strip() for w in args.workers.split(",") if w.strip()] if not workers: raise SystemExit("No workers specified") combined_manifest = Path(args.combined_manifest) case_rows = write_combined_manifest(combined_manifest) if len(case_rows) != 30: raise SystemExit(f"Expected 30 case packs, got {len(case_rows)}") completed_pairs = read_pairs(Path(args.completed_pairs)) if args.completed_pairs else set() target_pairs = read_pairs(Path(args.target_pairs)) if args.target_pairs else set() extensions = load_extensions() personas = list(iter_jsonl(PERSONA_FILE)) handles = {worker: (outdir / f"{worker}.jsonl").open("w", encoding="utf-8") for worker in workers} counts = {worker: 0 for worker in workers} report_counts = {} target_report_counts = {} skipped_completed = 0 skipped_not_target = 0 total = 0 written = 0 try: for case_row in case_rows: case_pack = load_case_pack(case_row["json_path"]) report_id = case_pack["case_pack_id"] fixed_prompt_prefix = { "common_instructions": COMMON_INSTRUCTIONS, "output_schema": OUTPUT_SCHEMA, "case_pack": case_pack, } prompt_hash = hashlib.sha256(dump(fixed_prompt_prefix).encode("utf-8")).hexdigest() report_counts[report_id] = 0 target_report_counts[report_id] = 0 for idx, rec in enumerate(personas, 1): pid = rec["persona_id"] pair = (report_id, pid) total += 1 report_counts[report_id] += 1 if target_pairs and pair not in target_pairs: skipped_not_target += 1 continue if pair in completed_pairs: skipped_completed += 1 continue ext = extensions.get(pid, {}) variable_input = { "assignment_group": f"{args.run_id}_r{case_pack['rank']:02d}_{idx:05d}", "persona": short_persona(rec), "extension_summary": { "derived_profile": ext.get("derived_profile"), "civic_and_administrative_orientation": ext.get("civic_and_administrative_orientation"), "evaluation_reception_profile": ext.get("evaluation_reception_profile"), "deliberation_style": ext.get("deliberation_style"), "preferred_explanation_format": ext.get("preferred_explanation_format"), "policy_interest_areas": ext.get("policy_interest_areas"), }, } row = { "persona_id": pid, "model_id": args.model, "dataset_version": args.run_id, "prompt_id": "resident_reaction_case_pack_v1_1", "prompt_hash": prompt_hash, "cache_key": f"{pid}:{prompt_hash}:{args.model}:{args.run_id}", "report_id": report_id, "condition_id": "original_case_pack_v1", "fixed_prompt_prefix": fixed_prompt_prefix, "variable_input": variable_input, } worker = workers[written % len(workers)] handles[worker].write(dump(row) + "\n") counts[worker] += 1 target_report_counts[report_id] += 1 written += 1 finally: for handle in handles.values(): handle.close() manifest = { "run_id": args.run_id, "outdir": str(outdir), "combined_manifest": str(combined_manifest), "case_pack_count": len(case_rows), "persona_count": len(personas), "total_eval30_pairs": total, "completed_pairs_loaded": len(completed_pairs), "target_pairs_loaded": len(target_pairs), "skipped_completed": skipped_completed, "skipped_not_target": skipped_not_target, "pending_rows_written": written, "workers": workers, "counts": counts, "all_report_counts": report_counts, "pending_report_counts": target_report_counts, } (outdir / "shard_manifest.json").write_text( json.dumps(manifest, ensure_ascii=False, indent=2) + "\n", encoding="utf-8", ) print(json.dumps(manifest, ensure_ascii=False, indent=2)) if __name__ == "__main__": main()