Files
hvm-docs/eval/trace.py
T
claude 7dbec5389d fix: default search_docs to dense; rerank opt-in
Post-prefix eval (22 queries, 2026-09-30): dense MRR 0.966 vs
bm25+rerank 0.827. Watchtower keeps live RERANK_URL, so rerank
is gated on RERANK_ENABLED (default off) rather than URL presence.
2026-09-29 22:45:21 -04:00

115 lines
3.9 KiB
Python

"""Page-level miss dump using this clone's Dense/BM25 retrievers.
python -m eval.trace --queries eval/queries.jsonl
"""
from __future__ import annotations
import argparse
import json
from pathlib import Path
from eval.run_eval import load_queries, p_at_1
def classify_top1_source(top1, dense_pages, bm25_pages) -> str:
if top1 is None:
return "neither"
in_d, in_b = top1 in dense_pages, top1 in bm25_pages
if in_d and in_b:
return "both"
if in_d:
return "dense_only"
if in_b:
return "bm25_only"
return "neither"
def first_ranks(pages: list[tuple[str, str]]) -> dict[str, int]:
out: dict[str, int] = {}
for i, (bid, pid) in enumerate(pages, start=1):
key = f"{bid}/{pid}"
if key not in out:
out[key] = i
return out
def render_misses(rows: list[dict]) -> str:
misses = [r for r in rows if not r.get("hit")]
if not misses:
return "# Eval misses\n\n_(none)_\n"
lines = [f"# Eval misses ({len(misses)})", ""]
for row in misses:
lines += [
f"## {row['query']}",
"",
f"- expected: `{row['expected']}`",
f"- top-5: `{row['ranked_pages'][:5]}`",
f"- top1_source: `{row['top1_source']}`",
"",
]
return "\n".join(lines)
def main() -> int:
p = argparse.ArgumentParser()
p.add_argument("--queries", type=Path, default=Path("eval/queries.jsonl"))
p.add_argument("--trace-out", type=Path, default=Path("eval/results/trace.jsonl"))
p.add_argument("--misses-out", type=Path, default=Path("eval/results/misses.md"))
args = p.parse_args()
if not args.queries.exists():
print(f"queries file not found: {args.queries}")
return 1
try:
import os
import chromadb
from chromadb.config import Settings
from rag.embeddings import embedding_function
from rag.bm25 import BM25Index
from eval.retrievers import BM25Retriever, DenseRetriever
product = os.environ.get("PRODUCT_NAME", "hvm")
root = Path(__file__).resolve().parent.parent
col = chromadb.PersistentClient(
path=str(root / "chroma"),
settings=Settings(anonymized_telemetry=False),
).get_collection(f"{product}_docs", embedding_function=embedding_function())
bm = BM25Index(str(root / "bm25" / f"{product}_docs.db"))
dense_r, bm25_r = DenseRetriever(col), BM25Retriever(bm)
except Exception as e:
args.trace_out.parent.mkdir(parents=True, exist_ok=True)
args.trace_out.write_text("")
args.misses_out.write_text("# Eval misses\n\nno index\n")
print(f"no index ({e}); wrote empty trace")
return 0
rows = []
for q in load_queries(args.queries):
expected = [(e["bundle_id"], e["page_id"]) for e in q["expected"]]
dense_pages = dense_r.retrieve(q["query"], k=50)
bm25_pages = bm25_r.retrieve(q["query"], k=50)
ranked = dense_pages or bm25_pages # HVM default retrieval is dense-first
top1 = ranked[0] if ranked else None
p1 = p_at_1(ranked, expected)
rows.append({
"query": q["query"],
"expected": [list(p) for p in expected],
"hit": bool(p1),
"p_at_1": int(p1),
"dense_rank": first_ranks(dense_pages),
"bm25_rank": first_ranks(bm25_pages),
"top1": list(top1) if top1 else None,
"top1_source": classify_top1_source(top1, set(dense_pages), set(bm25_pages)),
"ranked_pages": [list(p) for p in ranked],
})
args.trace_out.parent.mkdir(parents=True, exist_ok=True)
with open(args.trace_out, "w") as fh:
for row in rows:
fh.write(json.dumps(row) + "\n")
args.misses_out.write_text(render_misses(rows))
print(f"wrote {args.trace_out} ({len(rows)} queries)")
print(f"wrote {args.misses_out}")
return 0
if __name__ == "__main__":
raise SystemExit(main())