File size: 3,852 Bytes
81f30c1 | 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 | """Í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()
|