ZacharyZhang-NY
a6a1ef05f0
One stride over file order, truncated to the probe count, dropped the tail of the suite: a 300-probe sample held 128 code, 86 pdf, 36 screenshot, 26 image and 24 video queries and no conversation query at all, so every report built from it measured five of six modalities. The budget is now split by suite weight with a floor of one probe per modality, and a test pins that the frozen sample covers all six. The set's question type rides along on each query, and a second metric asks the question recall@10 cannot: when the answer lives in several dated memories, did the newest come back first? A knowledge-update question scores 1.000 on recall while the ranking still leads with the value the speaker has since replaced, and that is the failure the next task is meant to fix.
131 lines
4.8 KiB
Python
131 lines
4.8 KiB
Python
"""The frozen query set: what is asked, who asks it, and which probes to time."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections.abc import Sequence
|
|
from dataclasses import dataclass
|
|
|
|
from bench.sets_common import MODALITIES, SETS, read_jsonl
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class Query:
|
|
query_id: str
|
|
text: str
|
|
modality: str
|
|
gold: frozenset[str]
|
|
holder: str | None
|
|
answerable: bool
|
|
as_of: str | None
|
|
category: str | None = None
|
|
|
|
|
|
|
|
def load_queries(modalities: Sequence[str] = MODALITIES) -> list[Query]:
|
|
"""Every query the suite froze, in file order."""
|
|
queries: list[Query] = []
|
|
for modality in modalities:
|
|
path = SETS / f"{modality}.jsonl"
|
|
if not path.is_file():
|
|
raise SystemExit(f"{path}: missing set, run python -m bench.build_sets first")
|
|
for record in read_jsonl(path):
|
|
gold = frozenset(hit["source_id"] for hit in record["gold"])
|
|
answerable = not bool(record["unanswerable"])
|
|
if answerable != bool(gold):
|
|
raise SystemExit(
|
|
f"{path}: {record['query_id']} carries {len(gold)} gold hits but "
|
|
f"unanswerable={record['unanswerable']}"
|
|
)
|
|
queries.append(
|
|
Query(
|
|
query_id=record["query_id"],
|
|
text=record["question"],
|
|
modality=modality,
|
|
gold=gold,
|
|
holder=record["scope"]["holder"],
|
|
answerable=answerable,
|
|
as_of=record.get("as_of"),
|
|
category=(record.get("origin") or {}).get("category"),
|
|
)
|
|
)
|
|
if not queries:
|
|
raise SystemExit("no queries loaded")
|
|
return queries
|
|
|
|
|
|
|
|
def holder_index(modalities: Sequence[str] = MODALITIES) -> dict[str, str]:
|
|
"""source_id -> holder, from the corpus side of the suite."""
|
|
holders: dict[str, str] = {}
|
|
for modality in modalities:
|
|
path = SETS / "corpus" / f"{modality}.jsonl"
|
|
if not path.is_file():
|
|
raise SystemExit(f"{path}: missing corpus index, run python -m bench.build_sets first")
|
|
for record in read_jsonl(path):
|
|
holders[record["source_id"]] = record["holder"]
|
|
return holders
|
|
|
|
|
|
|
|
def source_times(modalities: Sequence[str] = MODALITIES) -> dict[str, str]:
|
|
"""source_id -> when that source says it happened, from the corpus side of the suite."""
|
|
times: dict[str, str] = {}
|
|
for modality in modalities:
|
|
path = SETS / "corpus" / f"{modality}.jsonl"
|
|
if not path.is_file():
|
|
raise SystemExit(f"{path}: missing corpus index, run python -m bench.build_sets first")
|
|
for record in read_jsonl(path):
|
|
observed = record.get("observed_at")
|
|
if observed:
|
|
times[record["source_id"]] = str(observed)
|
|
return times
|
|
|
|
|
|
|
|
SCALE_PROBES = 40
|
|
|
|
|
|
def _shares(pools: dict[str, list[Query]], count: int) -> dict[str, int]:
|
|
"""How many probes each modality gets: its share of the suite, never zero when there is room."""
|
|
names = sorted(pools)
|
|
total = sum(len(pools[name]) for name in names)
|
|
exact = {name: count * len(pools[name]) / total for name in names}
|
|
share = {name: int(exact[name]) for name in names}
|
|
short = count - sum(share.values())
|
|
for name in sorted(names, key=lambda name: (-(exact[name] - int(exact[name])), name))[:short]:
|
|
share[name] += 1
|
|
if count >= len(names):
|
|
for name in names:
|
|
while share[name] == 0:
|
|
donor = max((other for other in names if share[other] > 1), key=lambda other: share[other])
|
|
share[donor] -= 1
|
|
share[name] += 1
|
|
return {name: min(share[name], len(pools[name])) for name in names}
|
|
|
|
|
|
def scale_probes(queries: Sequence[Query], count: int = SCALE_PROBES) -> list[Query]:
|
|
"""Queries to time at every corpus size.
|
|
|
|
Latency does not need gold hits, and a truncated corpus keeps almost none of them,
|
|
so selecting probes by gold would silently drop the small sizes entirely. The budget is then
|
|
split across modalities by suite weight: one stride over file order truncated the tail and
|
|
timed 300 probes that held no conversation query at all.
|
|
"""
|
|
if count <= 0:
|
|
raise SystemExit("the scale curve needs at least one probe query")
|
|
pools: dict[str, list[Query]] = {}
|
|
for query in queries:
|
|
if query.answerable:
|
|
pools.setdefault(query.modality, []).append(query)
|
|
if not pools:
|
|
raise SystemExit("the scale curve needs answerable queries")
|
|
shares = _shares(pools, count)
|
|
chosen: list[Query] = []
|
|
for modality in sorted(pools):
|
|
if shares[modality] == 0:
|
|
continue
|
|
pool = pools[modality]
|
|
stride = max(1, len(pool) // shares[modality])
|
|
chosen.extend(pool[::stride][: shares[modality]])
|
|
return chosen
|