| """Índice ANN de `mathlib_full`: construye IVF_SQ (si hace falta) y mide cuánto se aparta de la |
| búsqueda exacta (recall@k del ANN respecto del top-k exacto), latencia p50/p95 de ambas y RSS. |
| |
| Consultas sintéticas sin cargar el embedder: normalize(e_i + e_j) para pares al azar de |
| declaraciones (caen entre dos documentos, no encima de uno). Uso: |
| python scripts/ann_recall.py [--build] [--n 200] [--out data/eval/results/ann_mathlib_full.json] |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import resource |
| import statistics |
| import time |
|
|
| import lancedb |
| import numpy as np |
|
|
| from ragsolver.rag.index import ANN_NPROBES, ANN_REFINE, build_ann_index, sql_quote |
|
|
|
|
| def main() -> None: |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--db-path", default="data/index/lancedb") |
| ap.add_argument("--table", default="mathlib_full") |
| ap.add_argument("--build", action="store_true") |
| ap.add_argument("--n", type=int, default=200) |
| ap.add_argument("--k", type=int, default=50) |
| ap.add_argument("--out", default="data/eval/results/ann_mathlib_full.json") |
| args = ap.parse_args() |
|
|
| t = lancedb.connect(args.db_path).open_table(args.table) |
| n_rows = t.count_rows() |
| if args.build: |
| t0 = time.time() |
| build_ann_index(t) |
| print(f"índice construido en {time.time() - t0:.0f} s", flush=True) |
| has_index = any("IVF" in str(i.index_type).upper() for i in t.list_indices()) |
|
|
| rng = np.random.default_rng(13) |
| all_ids = t.search().select(["chunk_id"]).limit(n_rows).to_arrow().column("chunk_id").to_pylist() |
| pick = [all_ids[i] for i in rng.choice(n_rows, size=2 * args.n, replace=False)] |
| rows = t.search().where("chunk_id IN (" + ", ".join(sql_quote(c) for c in pick) + ")").select(["chunk_id", "embedding"]).limit(len(pick)).to_list() |
| by_id = {r["chunk_id"]: r["embedding"] for r in rows} |
| emb = np.array([by_id[c] for c in pick], dtype="float32") |
| queries = emb[0::2] + emb[1::2] |
| queries /= np.linalg.norm(queries, axis=1, keepdims=True) |
|
|
| rec10, rec50, lat_exact, lat_ann = [], [], [], [] |
| for i, q in enumerate(queries): |
| t0 = time.time() |
| exact = [r["chunk_id"] for r in t.search(q, vector_column_name="embedding").bypass_vector_index().limit(args.k).select(["chunk_id"]).to_list()] |
| lat_exact.append(time.time() - t0) |
| t0 = time.time() |
| ann = [r["chunk_id"] for r in t.search(q, vector_column_name="embedding").nprobes(ANN_NPROBES).refine_factor(ANN_REFINE).limit(args.k).select(["chunk_id"]).to_list()] |
| lat_ann.append(time.time() - t0) |
| rec10.append(len(set(exact[:10]) & set(ann[:10])) / 10) |
| rec50.append(len(set(exact) & set(ann)) / len(exact)) |
| if (i + 1) % 50 == 0: |
| print(f"{i + 1}/{len(queries)} recall@10 {statistics.mean(rec10):.4f} recall@50 {statistics.mean(rec50):.4f}", flush=True) |
|
|
| p = lambda xs, q: round(sorted(xs)[int(q * (len(xs) - 1))], 3) |
| out = { |
| "table": args.table, "rows": n_rows, "ann_index": has_index, "index_type": "IVF_SQ", "num_partitions": 256, |
| "nprobes": ANN_NPROBES, "refine_factor": ANN_REFINE, "n_queries": len(queries), "k": args.k, |
| "ann_recall@10_vs_exact": round(statistics.mean(rec10), 4), "ann_recall@50_vs_exact": round(statistics.mean(rec50), 4), |
| "latency_exact_s": {"p50": p(lat_exact, 0.5), "p95": p(lat_exact, 0.95)}, "latency_ann_s": {"p50": p(lat_ann, 0.5), "p95": p(lat_ann, 0.95)}, |
| "maxrss_mb": resource.getrusage(resource.RUSAGE_SELF).ru_maxrss // 1024, |
| "note": "consultas sintéticas normalize(e_i + e_j); latencias medidas con la corrida eval_v1 en paralelo (CPU compartida)", |
| } |
| open(args.out, "w", encoding="utf-8").write(json.dumps(out, indent=2, ensure_ascii=False)) |
| print(json.dumps(out, ensure_ascii=False)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
