File size: 4,849 Bytes
33962ff | 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 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 | """Retrieval experiments: every (view x MRL dim) over the eval set, plus per-query ranks for the UI.
Writes data/experiments/results.json, plus results.<transcript engine>.json so runs over Whisper vs
Gemma transcripts can be compared side by side.
"""
import json
import time
from datetime import datetime, timezone
import numpy as np
from rich.console import Console
from rich.table import Table
from . import evalset
from .config import DIMS, EMBED_MODEL, EXPERIMENTS_DIR
from .embedder import get_model
from .index import get_index
from .segment import transcript_engine
console = Console()
KS = (1, 5, 10)
def metrics(ranks: list[int], video_top1: list[bool], neighbor_ranks: list[int]) -> dict:
r = np.array(ranks)
nb = np.array(neighbor_ranks)
out = {f"recall@{k}": float((r <= k).mean()) for k in KS}
out["mrr@10"] = float(np.where(r <= 10, 1.0 / r, 0.0).mean())
out["neighbor_recall@5"] = float((nb <= 5).mean())
out["video_hit@1"] = float(np.mean(video_top1))
out["n"] = len(ranks)
return out
def run() -> dict:
queries = evalset.load()
idx = get_index()
console.print(f"[bold]Running experiments[/bold]: {len(queries)} queries x {idx.available_views} x {DIMS}")
t0 = time.time()
qembs = get_model().encode([q["query"] for q in queries], prompt_name="SearchQuery",
batch_size=64, convert_to_numpy=True).astype(np.float32)
query_embed_ms = 1000 * (time.time() - t0) / len(queries)
seg_pos = {s["id"]: i for i, s in enumerate(idx.segments)}
per_query = [{**q, "ranks": {}} for q in queries]
grid = []
for view in idx.available_views:
for dim in DIMS:
ranks, video_top1, neighbor = [], [], []
t1 = time.time()
for qi, (q, qe) in enumerate(zip(queries, qembs)):
order, _ = idx.rank(qe, view, dim)
gold = seg_pos[q["segment_id"]]
rank = int(np.where(order == gold)[0][0]) + 1
gold_seg = idx.segments[gold]
# Lenient: a hit on the adjacent segment of the same video still lands the viewer in the right spot.
near = [r + 1 for r, i in enumerate(order[:10])
if idx.segments[i]["video_id"] == gold_seg["video_id"] and abs(idx.segments[i]["idx"] - gold_seg["idx"]) <= 1]
ranks.append(rank)
neighbor.append(near[0] if near else 10**6)
video_top1.append(idx.segments[order[0]]["video_id"] == q["video_id"])
per_query[qi]["ranks"][f"{view}@{dim}"] = rank
search_ms = 1000 * (time.time() - t1) / len(queries)
for qtype in ("all", "spoken", "visual"):
sel = [i for i, q in enumerate(queries) if qtype == "all" or q["type"] == qtype]
if not sel:
continue
grid.append({
"view": view, "dim": dim, "query_type": qtype,
**metrics([ranks[i] for i in sel], [video_top1[i] for i in sel], [neighbor[i] for i in sel]),
"search_ms": round(search_ms, 3),
})
views_count = len([v for v in idx.available_views if v != "hybrid"])
result = {
"created_at": datetime.now(timezone.utc).isoformat(timespec="seconds"),
"model": EMBED_MODEL,
"transcripts": transcript_engine(),
"num_segments": len(idx.segments),
"num_queries": len(queries),
"query_types": {t: sum(q["type"] == t for q in queries) for t in ("spoken", "visual")},
"query_embed_ms": round(query_embed_ms, 2),
"dims": DIMS,
"views": idx.available_views,
"storage": [{"dim": d, "bytes_per_vector": d * 4, "bytes_per_view": idx.storage_bytes(d),
"bytes_all_views": idx.storage_bytes(d) * views_count} for d in DIMS],
"grid": grid,
"per_query": per_query,
}
EXPERIMENTS_DIR.mkdir(parents=True, exist_ok=True)
(EXPERIMENTS_DIR / "results.json").write_text(json.dumps(result, indent=1))
(EXPERIMENTS_DIR / f"results.{result['transcripts']}.json").write_text(json.dumps(result, indent=1))
_print(result)
return result
def _print(result: dict) -> None:
for qtype in ("all", "spoken", "visual"):
t = Table(title=f"recall@5 / MRR@10 - {qtype} queries")
t.add_column("view")
for d in result["dims"]:
t.add_column(f"{d}d", justify="right")
for view in result["views"]:
cells = []
for d in result["dims"]:
row = next((g for g in result["grid"] if g["view"] == view and g["dim"] == d and g["query_type"] == qtype), None)
cells.append(f"{row['recall@5']:.2f} / {row['mrr@10']:.2f}" if row else "-")
t.add_row(view, *cells)
console.print(t)
|