TeX 59%
Python 40.8%
1#!/usr/bin/env python32"""Step 05 — parse all batch results, build per-image and property-level tables for every tier and image budget K,3and the baseline property tables B0–B4 (per tier). Writes:4 data/processed/images_<tier>.parquet per-image extraction rows (ok + errors)5 data/processed/text_<tier>.parquet per-listing text extraction rows6 data/processed/props_<tier>_K<k>.parquet aggregated visual attributes (K = 1,2,4,8,10,12 where available)7 data/processed/levels_<tier>.parquet long table: level ∈ {B0..B4} × uid with cost-model inputs8 results/raw/aggregate_summary.json9Also handles the supplementary "sold" sample (suffix _sold) when present.10"""11import json12import sys13from pathlib import Path1415import pandas as pd1617sys.path.insert(0, str(Path(__file__).resolve().parents[1]))18from src.extraction.aggregate import parse_results, aggregate_images, build_levels # noqa: E4021920ROOT = Path(__file__).resolve().parents[1]21RAW = ROOT / "results" / "raw" / "batches"22OUT = ROOT / "data" / "processed"23KS = {"haiku": [1, 2, 4, 8, 10, 12], "sonnet": [1, 2, 4, 8, 10], "opus": [1, 2, 4, 8, 10]}242526def main(suffix: str = "") -> None:27 frame_file = "sample_sold.parquet" if suffix else "sample.parquet"28 smp = pd.read_parquet(OUT / frame_file)29 summary = {}30 texts = {}31 for tier in ["haiku", "sonnet", "opus"]:32 f = RAW / f"txt_{tier}{suffix}.results.jsonl"33 if f.exists():34 t = parse_results(f"txt_{tier}{suffix}")35 t.to_parquet(OUT / f"text_{tier}{suffix}.parquet", index=False)36 texts[tier] = t37 summary[f"text_{tier}"] = {"n": int(len(t)), "ok": int(t.ok.sum())}38 for tier in ["haiku", "sonnet", "opus"]:39 f = RAW / f"img_{tier}{suffix}.results.jsonl"40 if not f.exists():41 continue42 img = parse_results(f"img_{tier}{suffix}")43 img.to_parquet(OUT / f"images_{tier}{suffix}.parquet", index=False)44 summary[f"img_{tier}"] = {"n": int(len(img)), "ok": int(img.ok.sum()), "listings": int(img.uid.nunique()),45 "mean_in_tokens": float(img.input_tokens.mean()), "mean_out_tokens": float(img.output_tokens.mean())}46 props_full = None47 for k in KS[tier]:48 p = aggregate_images(img, K=k)49 p.to_parquet(OUT / f"props_{tier}{suffix}_K{k}.parquet", index=False)50 props_full = p51 txt = texts.get(tier) if tier in texts else texts.get("sonnet")52 levels = build_levels(smp, props_full, txt)53 long = pd.concat([df.assign(level=lvl).reset_index() for lvl, df in levels.items()], ignore_index=True)54 long.to_parquet(OUT / f"levels_{tier}{suffix}.parquet", index=False)55 summary[f"levels_{tier}"] = {lvl: int(len(df)) for lvl, df in levels.items()}56 json.dump(summary, open(ROOT / "results" / "raw" / f"aggregate_summary{suffix}.json", "w"), indent=2)57 print(json.dumps(summary, indent=1))585960if __name__ == "__main__":61 main(sys.argv[1] if len(sys.argv) > 1 else "")62