File size: 3,176 Bytes
e6f80f3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 | """A fixed test subset for the LLM comparisons: N rooms of test.jsonl, stratified by source in proportion to the test
set (largest remainder; a source too small for one room gets none), each source's rooms drawn with one seeded RNG.
Rooms over the model capacity (more than --max-objects objects) are left out before drawing, so every method can
answer every room. Writes rows.jsonl (test order) and MANIFEST.json (seed, counts, sha256, scene ids).
python fastfill/v2/ops/make_test300.py --test data/v4fixed/test.jsonl --output data/test300 [--rows 300 --seed 0]
"""
from __future__ import annotations
import argparse
from collections import defaultdict
import hashlib
import json
from pathlib import Path
import random
def allocate(sizes: dict, total: int) -> dict:
"""Largest-remainder apportionment of ``total`` over ``sizes`` (ties broken by name)."""
whole = sum(sizes.values())
quota = {k: total * v / whole for k, v in sizes.items()}
share = {k: int(q) for k, q in quota.items()}
for k in sorted(quota, key=lambda k: (-(quota[k] - share[k]), k))[:total - sum(share.values())]:
share[k] += 1
return share
def main(argv=None):
p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
p.add_argument("--test", type=Path, required=True)
p.add_argument("--output", type=Path, required=True)
p.add_argument("--rows", type=int, default=300)
p.add_argument("--seed", type=int, default=0)
p.add_argument("--max-objects", type=int, default=256)
a = p.parse_args(argv)
lines = [l for l in a.test.read_text().splitlines() if l.strip()]
by_source, over = defaultdict(list), 0
for index, line in enumerate(lines):
row = json.loads(line)
if len(row["condition"]["objects"]) > a.max_objects:
over += 1
continue
by_source[row["provenance"]["source"]].append(index)
share, rng, chosen = allocate({k: len(v) for k, v in by_source.items()}, a.rows), random.Random(a.seed), []
for source in sorted(by_source):
chosen += rng.sample(by_source[source], share[source])
chosen.sort()
a.output.mkdir(parents=True, exist_ok=True)
rows = "".join(lines[i] + "\n" for i in chosen)
(a.output / "rows.jsonl").write_text(rows)
manifest = {"rows": len(chosen), "seed": a.seed, "max_objects": a.max_objects, "over_capacity_excluded": over,
"test_file": a.test.name, "test_rows": len(lines), "test_sha256": hashlib.sha256(a.test.read_bytes()).hexdigest(),
"rows_sha256": hashlib.sha256(rows.encode()).hexdigest(),
"per_source": {k: {"test": len(by_source[k]), "chosen": share[k]} for k in sorted(by_source)},
"test_line_indices": chosen,
"scene_ids": [json.loads(lines[i])["provenance"].get("scene_id") for i in chosen]}
(a.output / "MANIFEST.json").write_text(json.dumps(manifest, indent=1) + "\n")
print(json.dumps({k: manifest[k] for k in ("rows", "over_capacity_excluded", "rows_sha256")}),
{k: v["chosen"] for k, v in manifest["per_source"].items()})
if __name__ == "__main__":
main()
|