| |
| """Independent MaxSim on a dump (raw) or HAC index. |
| |
| Does not re-encode. Reports NDCG / Recall / MAP / MRR at @1/@5/@10. |
| Search timing is query-already-encoded MaxSim only — not product QPS. |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import os |
| import sys |
| import time |
| from collections import defaultdict |
| from pathlib import Path |
|
|
| sys.path.insert(0, str(Path(__file__).resolve().parent)) |
| sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "shared")) |
|
|
| import torch |
|
|
| from common import ( |
| BUDGET, |
| CEILING, |
| DUMP_PROTOCOL, |
| HEAD_DIM, |
| MVT, |
| REQUIRED_KS, |
| REQUIRED_METRICS, |
| SOUP_NAME, |
| WEIGHT, |
| atomic_write_json, |
| dump_ok, |
| gib_per_million_pages, |
| has_complete_metrics, |
| load_qrels, |
| padded_from_pack, |
| w_tag, |
| ) |
|
|
| import eval as evie_eval |
| from paths import forbid_venv_path |
|
|
|
|
| DS_LABEL = evie_eval.DS_LABEL |
| METRIC_NAMES = evie_eval.METRIC_NAMES |
|
|
|
|
| def load_task_list(dump_root: Path) -> list[dict]: |
| tasks_dir = dump_root / "tasks" |
| recs = [] |
| for path in sorted(tasks_dir.glob("*.json")): |
| if dump_ok(path, dump_root): |
| recs.append(json.loads(path.read_text(encoding="utf-8"))) |
| if not recs: |
| raise SystemExit(f"no complete dump tasks under {tasks_dir}") |
| return recs |
|
|
|
|
| def load_queries(dump_root: Path, rec: dict) -> torch.Tensor: |
| blob = torch.load(dump_root / rec["queries"], map_location="cpu", weights_only=False) |
| return padded_from_pack(blob["packed"], blob["offsets"]) |
|
|
|
|
| def load_raw_docs(dump_root: Path, rec: dict) -> torch.Tensor: |
| blob = torch.load(dump_root / rec["docs"], map_location="cpu", weights_only=False) |
| return padded_from_pack(blob["packed"], blob["offsets"]) |
|
|
|
|
| def load_clustered_docs(cluster_root: Path, rec: dict, tag: str, budget: int) -> tuple[torch.Tensor, dict]: |
| path = Path(cluster_root) / tag / "docs" / f"{rec['corpus_id']}.pt" |
| blob = torch.load(path, map_location="cpu", weights_only=False) |
| padded = blob["padded"] |
| if padded.ndim != 3 or int(padded.shape[-1]) != HEAD_DIM: |
| raise RuntimeError(f"bad clustered tensor {path}: {tuple(padded.shape)}") |
| if int(padded.shape[1]) != int(budget): |
| raise RuntimeError(f"clustered K {padded.shape[1]} != budget {budget} at {path}") |
| return padded, blob.get("stats") or {} |
|
|
|
|
| def score_task(q_emb, d_emb, relevant, graded, device, chunk_q, chunk_d): |
| q = q_emb.to(device) |
| t0 = time.perf_counter() |
| scores = evie_eval.chunked_maxsim(q, d_emb, chunk_q=chunk_q, chunk_d=chunk_d) |
| dt = time.perf_counter() - t0 |
| metrics = evie_eval.metrics_from_scores(scores, relevant, graded, list(REQUIRED_KS)) |
| n_q = int(q_emb.shape[0]) |
| return metrics, { |
| "n_docs": int(d_emb.shape[0]), |
| "n_queries": n_q, |
| "doc_tokens": int(d_emb.shape[1]), |
| "query_tokens_padded": int(q_emb.shape[1]), |
| "search_sec": round(dt, 4), |
| "search_ms_per_query": round(1000.0 * dt / max(n_q, 1), 4), |
| "encode_ms_per_query": None, |
| "note": "pure search; query vectors already dumped; not product QPS", |
| "chunk_q": chunk_q, |
| "chunk_d": chunk_d, |
| "device": str(device), |
| } |
|
|
|
|
| def write_summary(out_dir: Path, recs: list[dict], expected: list[dict], args, mode: str, w): |
| by_ds = defaultdict(list) |
| failed = [r for r in recs if "error" in r] |
| ok = [r for r in recs if "metrics" in r and has_complete_metrics(r["metrics"])] |
| for rec in ok: |
| by_ds[rec["dataset"]].append(rec) |
| expected_by_ds = defaultdict(int) |
| for rec in expected: |
| expected_by_ds[rec["dataset"]] += 1 |
| expected_tags = {rec["task"] for rec in expected} |
| missing = sorted(expected_tags - {r["task"] for r in ok}) |
|
|
| def avg_metric(rows, name): |
| vals = [r["metrics"].get(name) for r in rows] |
| vals = [v for v in vals if v is not None] |
| return (sum(vals) / len(vals)) if vals else None |
|
|
| ks = list(REQUIRED_KS) |
| per_ds = {} |
| for ds, rows in by_ds.items(): |
| entry = {f"{m}@{k}": avg_metric(rows, f"{m}@{k}") for m in METRIC_NAMES for k in ks} |
| entry["n_subsets"] = len(rows) |
| entry["n_subsets_expected"] = expected_by_ds.get(ds, len(rows)) |
| entry["n_queries"] = sum(r["info"]["n_queries"] for r in rows) |
| entry["n_docs"] = sum(r["info"]["n_docs"] for r in rows) |
| search = [r["info"]["search_ms_per_query"] for r in rows if r["info"].get("search_ms_per_query") is not None] |
| entry["search_ms_per_query_mean_unweighted"] = ( |
| sum(search) / len(search) if search else None |
| ) |
| per_ds[ds] = entry |
|
|
| headline = { |
| "V1 nDCG@10": per_ds.get("eval", {}).get("ndcg@10"), |
| "V2 nDCG@10": per_ds.get("eval_v2", {}).get("ndcg@10"), |
| "V3 nDCG@10": per_ds.get("eval_v3", {}).get("ndcg@10"), |
| "JinaVDR nDCG@10": per_ds.get("JinaVDR", {}).get("ndcg@10"), |
| } |
| four = [headline[k] for k in ("V1 nDCG@10", "V2 nDCG@10", "V3 nDCG@10", "JinaVDR nDCG@10")] |
| averages = { |
| "avg4_ndcg@10": (sum(four) / 4 if all(v is not None for v in four) else None), |
| } |
| complete = not failed and not missing |
| vec_per_page = 1.0 |
| if mode == "k": |
| vec_per_page = float(args.budget) |
| storage_note = f"independent K{int(args.budget)} candidate layer, no context, no raw rerank" |
| else: |
| vec_per_page = None |
| storage_note = "raw full sequence from dump; not a compressed index" |
| summary = { |
| "status": "complete" if complete else "partial", |
| "run_name": args.run_name, |
| "mode": mode, |
| "position_weight": None if mode == "raw" else float(w), |
| "budget": None if mode == "raw" else int(args.budget), |
| "export_context": False, |
| "head_dim": HEAD_DIM, |
| "max_visual_tokens": MVT, |
| "bidirectional_attention": "on", |
| "weights": SOUP_NAME, |
| "alpha": 0.5, |
| "dump_protocol": DUMP_PROTOCOL, |
| "eval_protocol": evie_eval.EVAL_PROTOCOL, |
| "ks": ks, |
| "expected_tasks": len(expected), |
| "completed_tasks": len(ok), |
| "n_failed": len(failed), |
| "failed_tasks": failed, |
| "missing_tasks": missing, |
| "headline": headline, |
| "averages": averages, |
| "per_dataset": per_ds, |
| "per_task": ok, |
| "ceiling": CEILING, |
| "storage": { |
| "vectors_per_page": vec_per_page, |
| "dim": HEAD_DIM, |
| "bytes_per_vec": 2, |
| "gib_per_million_pages": ( |
| None if vec_per_page is None else round(gib_per_million_pages(vec_per_page), 4) |
| ), |
| "context": False, |
| "two_stage": False, |
| "note": storage_note, |
| }, |
| "search": { |
| "kind": "independent_pure_search", |
| "includes_query_encode": False, |
| "not_product_qps": True, |
| "two_stage": False, |
| }, |
| } |
| atomic_write_json(out_dir / "summary.json", summary) |
| lines = [ |
| f"# {args.run_name} {mode}" |
| + (f" K{int(args.budget)} {w_tag(w)}" if mode == "k" else ""), |
| "", |
| f"status `{summary['status']}` · {summary['completed_tasks']}/{summary['expected_tasks']}", |
| "", |
| "Independent retrieval. Queries are already dumped; ms below is MaxSim only, not product QPS.", |
| "", |
| "| board | nDCG@1 | nDCG@5 | nDCG@10 | Recall@1 | Recall@5 | Recall@10 | MAP@1 | MAP@5 | MAP@10 | MRR@1 | MRR@5 | MRR@10 |", |
| "| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: |", |
| ] |
|
|
| def pct(v): |
| return "—" if v is None else f"{v * 100:.2f}" |
|
|
| for ds, label in DS_LABEL.items(): |
| e = per_ds.get(ds) |
| if not e: |
| continue |
| cells = " | ".join( |
| pct(e.get(name)) |
| for name in ( |
| "ndcg@1", |
| "ndcg@5", |
| "ndcg@10", |
| "recall@1", |
| "recall@5", |
| "recall@10", |
| "map@1", |
| "map@5", |
| "map@10", |
| "mrr@1", |
| "mrr@5", |
| "mrr@10", |
| ) |
| ) |
| lines.append(f"| {label} | {cells} |") |
| lines += [ |
| "", |
| f"uncompressed d{HEAD_DIM} eval V3 nDCG@10 = {CEILING['V3 nDCG@10']} " |
| f"(source `{CEILING['source']}`)", |
| "", |
| ] |
| (out_dir / "summary.md").write_text("\n".join(lines) + "\n", encoding="utf-8") |
| print( |
| f"[eval] {summary['status']} {summary['completed_tasks']}/{summary['expected_tasks']} " |
| f"V3 nDCG@10={headline.get('V3 nDCG@10')}", |
| flush=True, |
| ) |
| if not complete: |
| raise SystemExit(2) |
|
|
|
|
| def parse_args(): |
| ap = argparse.ArgumentParser(description=__doc__) |
| ap.add_argument("--dump-dir", required=True) |
| ap.add_argument("--output-dir", required=True) |
| ap.add_argument("--cluster-root", default="") |
| ap.add_argument("--mode", choices=("raw", "k"), required=True) |
| ap.add_argument("--w", type=float, default=WEIGHT) |
| ap.add_argument("--budget", type=int, default=BUDGET) |
| ap.add_argument("--run-name", default="evie-compress") |
| ap.add_argument("--chunk-q", type=int, default=64) |
| ap.add_argument("--chunk-d", type=int, default=256) |
| ap.add_argument("--resume", action="store_true") |
| return ap.parse_args() |
|
|
|
|
| def main(): |
| args = parse_args() |
| if args.mode == "k": |
| if args.w is None: |
| raise SystemExit("k mode requires --w") |
| if not args.cluster_root: |
| raise SystemExit("k mode requires --cluster-root") |
| tag = w_tag(args.w) |
| else: |
| tag = "raw" |
| args.w = None |
|
|
| rank, world, local = evie_eval.setup_ddp() |
| device = torch.device(f"cuda:{local}" if torch.cuda.is_available() else "cpu") |
| dump_root = forbid_venv_path(args.dump_dir, "dump-dir") |
| out_dir = forbid_venv_path(args.output_dir, "output-dir") |
| out_dir.mkdir(parents=True, exist_ok=True) |
| expected = load_task_list(dump_root) |
|
|
| cached = {} |
| if args.resume: |
| cached = evie_eval.load_cached_task_map(out_dir, protocol=evie_eval.EVAL_PROTOCOL) |
| cached = {k: v for k, v in cached.items() if has_complete_metrics(v.get("metrics"))} |
|
|
| pending = [rec for rec in expected if rec["task"] not in cached] |
| if world > 1: |
| |
| pending = sorted(pending, key=lambda r: -(int(r["n_docs"]) * int(r["n_queries"]))) |
| pending = pending[rank::world] |
| print(f"[eval rank{rank}] {args.mode} {tag} pending {len(pending)}/{len(expected)} on {device}", flush=True) |
|
|
| results = [] |
| for rec in pending: |
| tag_task = rec["task"] |
| t0 = time.time() |
| try: |
| q_emb = load_queries(dump_root, rec) |
| if args.mode == "raw": |
| d_emb = load_raw_docs(dump_root, rec) |
| else: |
| d_emb, _stats = load_clustered_docs(Path(args.cluster_root), rec, tag, args.budget) |
| relevant, graded = load_qrels(dump_root / rec["qrels"]) |
| if len(relevant) != q_emb.shape[0]: |
| raise RuntimeError( |
| f"{tag_task}: qrels {len(relevant)} != queries {q_emb.shape[0]}" |
| ) |
| if d_emb.shape[0] != rec["n_docs"]: |
| raise RuntimeError( |
| f"{tag_task}: docs {d_emb.shape[0]} != dump n_docs {rec['n_docs']}" |
| ) |
| metrics, info = score_task( |
| q_emb, d_emb, relevant, graded, device, args.chunk_q, args.chunk_d |
| ) |
| row = { |
| "task": tag_task, |
| "dataset": rec["dataset"], |
| "subset": rec["subset"], |
| "lang": rec.get("lang"), |
| "fmt": rec["fmt"], |
| "metrics": metrics, |
| "info": info, |
| "seconds": round(time.time() - t0, 1), |
| "eval_protocol": evie_eval.EVAL_PROTOCOL, |
| "dump_protocol": DUMP_PROTOCOL, |
| "run_name": args.run_name, |
| "mode": args.mode, |
| "position_weight": args.w, |
| } |
| evie_eval.save_task_result(out_dir, row, run_name=args.run_name, contract_id=DUMP_PROTOCOL) |
| results.append(row) |
| print( |
| f"[eval rank{rank}] {tag_task}: ndcg@10={metrics['ndcg@10']*100:.2f} " |
| f"{info['n_queries']}q/{info['n_docs']}d search={info['search_ms_per_query']:.2f}ms/q", |
| flush=True, |
| ) |
| except Exception as exc: |
| import traceback |
|
|
| traceback.print_exc() |
| results.append({"task": tag_task, "dataset": rec["dataset"], "error": repr(exc)}) |
| print(f"[eval rank{rank}] {tag_task} FAILED: {exc!r}", flush=True) |
|
|
| atomic_write_json(out_dir / f"rank_{rank}.json", results) |
| if world > 1: |
| torch.distributed.destroy_process_group() |
| if rank != 0: |
| return |
|
|
| deadline = time.monotonic() + int(os.environ.get("EVAL_FINALIZE_TIMEOUT_S", "21600")) |
| missing_ranks = list(range(world)) |
| last = 0.0 |
| while missing_ranks: |
| missing_ranks = [] |
| for r in range(world): |
| p = out_dir / f"rank_{r}.json" |
| try: |
| json.loads(p.read_text(encoding="utf-8")) |
| except Exception: |
| missing_ranks.append(r) |
| if not missing_ranks: |
| break |
| now = time.monotonic() |
| if now >= deadline: |
| print(f"[eval] finalize timeout; missing {missing_ranks}", flush=True) |
| break |
| if now - last >= 60: |
| print(f"[eval] waiting for ranks {missing_ranks}", flush=True) |
| last = now |
| time.sleep(2) |
|
|
| merged = evie_eval.load_cached_task_map(out_dir, protocol=evie_eval.EVAL_PROTOCOL) |
| rows = [] |
| failed = [] |
| for r in range(world): |
| p = out_dir / f"rank_{r}.json" |
| if not p.exists(): |
| continue |
| for rec in json.loads(p.read_text(encoding="utf-8")): |
| if "error" in rec: |
| failed.append(rec) |
| elif rec.get("task") in merged: |
| rows.append(merged[rec["task"]]) |
| |
| rows = [merged[rec["task"]] for rec in expected if rec["task"] in merged] |
| for miss in failed: |
| if miss["task"] not in merged: |
| rows.append(miss) |
| write_summary(out_dir, rows, expected, args, args.mode, args.w) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
