| """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] |
| |
| 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) |
|
|
