1"""Phase C — assign table dumps to extraction batches (balanced by dump size) and write the batch manifests.2Usage: python scripts/extraction/make_batches.py --n 8 [--only screened] -> data/extracted/batches/batch_<i>.txt (list of record_ids)."""3import argparse, json4from pathlib import Path5import pandas as pd67ROOT = Path(__file__).resolve().parents[2]; EX = ROOT / "data" / "extracted"8910def main():11 ap = argparse.ArgumentParser(); ap.add_argument("--n", type=int, default=8); ap.add_argument("--max_chars", type=int, default=60000); ap.add_argument("--start", type=int, default=1); a = ap.parse_args()12 idx = pd.read_csv(EX / "pdf_index.csv"); idx = idx[(~idx.scanned) & (idx.n_candidate_pages > 0)]13 done = {p.stem for p in (EX / "json").glob("*.json")} if (EX / "json").exists() else set()14 assigned = set()15 for b in (EX / "batches").glob("batch_*.txt"):16 assigned |= set(b.read_text().split())17 idx = idx[~idx.record_id.astype(str).isin(done | assigned)]18 # screening decisions (if available): keep includes/unclear19 dec_files = sorted((ROOT / "data" / "metadata" / "screening_batches").glob("decisions_*.csv"))20 if dec_files:21 dec = pd.concat(pd.read_csv(f) for f in dec_files)22 keep = set(dec[dec.reason_code.eq("ELIGIBLE_HEDONIC") | dec.decision.eq("unclear")].record_id.astype(str))23 idx = idx[idx.record_id.astype(str).isin(keep) | ~idx.record_id.astype(str).isin(set(dec.record_id.astype(str)))]24 idx = idx.sort_values("dump_chars", ascending=False)25 batches = [[] for _ in range(a.n)]; sizes = [0] * a.n26 for r in idx.itertuples():27 j = sizes.index(min(sizes)); batches[j].append(r.record_id); sizes[j] += min(r.dump_chars, a.max_chars)28 (EX / "batches").mkdir(exist_ok=True)29 for i, b in enumerate(batches, a.start):30 (EX / "batches" / f"batch_{i}.txt").write_text("\n".join(map(str, b)))31 print(len(idx), "dumps in", a.n, "batches; chars per batch:", sizes)323334if __name__ == "__main__":35 main()36