Port docs-mcp-template upgrades (eval, citations, chunking, prefixes) (#16)
Co-authored-by: claude <[email protected]>
This commit was merged in pull request #16.
This commit is contained in:
+156
@@ -0,0 +1,156 @@
|
||||
"""Paired permutation test between two eval JSONL sidecars.
|
||||
|
||||
Compares per-query scores from two `eval.run_eval` sidecar files so
|
||||
"P@1 went 0.88 → 0.91 on 25 queries" is not treated as a win.
|
||||
|
||||
python -m eval.pvalue \\
|
||||
--a eval/results/baseline.jsonl \\
|
||||
--b eval/results/new.jsonl \\
|
||||
--metric rr
|
||||
|
||||
Exit 0 even when the difference is not significant — this is a report,
|
||||
not a gate. No third-party deps; `random.Random(seed)` is enough.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import random
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
def load_sidecar(path: Path) -> list[dict]:
|
||||
rows: list[dict] = []
|
||||
with open(path) as fh:
|
||||
for line in fh:
|
||||
line = line.strip()
|
||||
if line:
|
||||
rows.append(json.loads(line))
|
||||
return rows
|
||||
|
||||
|
||||
def paired_permutation(
|
||||
a: list[float],
|
||||
b: list[float],
|
||||
n_resamples: int = 10000,
|
||||
seed: int = 0,
|
||||
) -> dict:
|
||||
"""Two-sided paired permutation test on per-query scores.
|
||||
|
||||
Null: each pair is exchangeable (randomly flipping the sign of
|
||||
A_i - B_i). p_value is the fraction of permutations whose
|
||||
|mean diff| is at least as large as the observed |mean(A-B)|.
|
||||
"""
|
||||
if len(a) != len(b):
|
||||
raise ValueError(f"paired lengths differ: {len(a)} vs {len(b)}")
|
||||
if not a:
|
||||
raise ValueError("no paired queries to compare")
|
||||
diffs = [x - y for x, y in zip(a, b)]
|
||||
n = len(diffs)
|
||||
observed = sum(diffs) / n
|
||||
abs_obs = abs(observed)
|
||||
rng = random.Random(seed)
|
||||
extreme = 0
|
||||
for _ in range(n_resamples):
|
||||
total = 0.0
|
||||
for d in diffs:
|
||||
total += d if rng.random() < 0.5 else -d
|
||||
if abs(total / n) >= abs_obs - 1e-15:
|
||||
extreme += 1
|
||||
p_value = extreme / n_resamples
|
||||
return {
|
||||
"A_mean": sum(a) / n,
|
||||
"B_mean": sum(b) / n,
|
||||
"Diff(A-B)": observed,
|
||||
"p_value": p_value,
|
||||
"significant": p_value < 0.05,
|
||||
"n": n,
|
||||
"n_resamples": n_resamples,
|
||||
}
|
||||
|
||||
|
||||
def _index(rows: list[dict], metric: str) -> dict[tuple[str, str], float]:
|
||||
"""Map (retriever, query) -> score."""
|
||||
out: dict[tuple[str, str], float] = {}
|
||||
for row in rows:
|
||||
retriever = str(row.get("retriever") or "")
|
||||
query = str(row.get("query") or "")
|
||||
if metric == "p_at_1":
|
||||
score = float(row.get("p_at_1") or 0)
|
||||
else:
|
||||
score = float(row.get("rr") or 0)
|
||||
out[(retriever, query)] = score
|
||||
return out
|
||||
|
||||
|
||||
def compare(
|
||||
rows_a: list[dict],
|
||||
rows_b: list[dict],
|
||||
metric: str = "rr",
|
||||
retriever: str | None = None,
|
||||
n_resamples: int = 10000,
|
||||
seed: int = 0,
|
||||
) -> list[dict]:
|
||||
"""Join on (retriever, query). One result dict per shared retriever."""
|
||||
ia, ib = _index(rows_a, metric), _index(rows_b, metric)
|
||||
retrievers = sorted({r for r, _ in ia} & {r for r, _ in ib})
|
||||
if retriever:
|
||||
retrievers = [r for r in retrievers if r == retriever]
|
||||
if not retrievers:
|
||||
raise ValueError(f"retriever {retriever!r} not in both sidecars")
|
||||
reports = []
|
||||
for name in retrievers:
|
||||
queries = sorted({q for r, q in ia if r == name} & {q for r, q in ib if r == name})
|
||||
if not queries:
|
||||
continue
|
||||
a_scores = [ia[(name, q)] for q in queries]
|
||||
b_scores = [ib[(name, q)] for q in queries]
|
||||
report = paired_permutation(a_scores, b_scores, n_resamples=n_resamples, seed=seed)
|
||||
report["retriever"] = name
|
||||
report["metric"] = metric
|
||||
reports.append(report)
|
||||
if not reports:
|
||||
raise ValueError("no overlapping (retriever, query) pairs")
|
||||
return reports
|
||||
|
||||
|
||||
def render(reports: list[dict]) -> str:
|
||||
lines = ["# Permutation test", ""]
|
||||
for r in reports:
|
||||
sig = "yes" if r["significant"] else "no"
|
||||
lines += [
|
||||
f"## `{r['retriever']}` ({r['metric']}, n={r['n']})",
|
||||
"",
|
||||
f"- A_mean: `{r['A_mean']:.4f}`",
|
||||
f"- B_mean: `{r['B_mean']:.4f}`",
|
||||
f"- Diff(A-B): `{r['Diff(A-B)']:.4f}`",
|
||||
f"- p_value: `{r['p_value']:.4f}` ({r['n_resamples']} resamples)",
|
||||
f"- significant (p < 0.05): **{sig}**",
|
||||
"",
|
||||
]
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
def main() -> int:
|
||||
p = argparse.ArgumentParser(description="Paired permutation test on two eval JSONL sidecars.")
|
||||
p.add_argument("--a", type=Path, required=True, help="sidecar JSONL (system A)")
|
||||
p.add_argument("--b", type=Path, required=True, help="sidecar JSONL (system B)")
|
||||
p.add_argument("--metric", choices=("rr", "p_at_1"), default="rr")
|
||||
p.add_argument("--retriever", default=None, help="restrict to one retriever name")
|
||||
p.add_argument("--n-resamples", type=int, default=10000)
|
||||
p.add_argument("--seed", type=int, default=0)
|
||||
args = p.parse_args()
|
||||
reports = compare(
|
||||
load_sidecar(args.a),
|
||||
load_sidecar(args.b),
|
||||
metric=args.metric,
|
||||
retriever=args.retriever,
|
||||
n_resamples=args.n_resamples,
|
||||
seed=args.seed,
|
||||
)
|
||||
print(render(reports), end="")
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
+3
-1
@@ -51,7 +51,9 @@ class DenseRetriever:
|
||||
self.pool = pool
|
||||
|
||||
def retrieve(self, query: str, k: int = 10) -> list[tuple[str, str]]:
|
||||
res = self.col.query(query_texts=[query], n_results=self.pool)
|
||||
from rag.embeddings import EMBED_QUERY_PREFIX, embed_texts
|
||||
qvec = embed_texts([query], prefix=EMBED_QUERY_PREFIX)[0]
|
||||
res = self.col.query(query_embeddings=[qvec], n_results=self.pool)
|
||||
ids = (res.get("ids") or [[]])[0]
|
||||
return _collapse_to_pages(ids, k)
|
||||
|
||||
|
||||
+37
-7
@@ -34,6 +34,12 @@ def load_queries(path: Path) -> list[dict]:
|
||||
return [json.loads(line) for line in fh if line.strip()]
|
||||
|
||||
|
||||
def p_at_1(retrieved: list[tuple[str, str]], expected: list[tuple[str, str]]) -> float:
|
||||
if not retrieved or not expected:
|
||||
return 0.0
|
||||
return 1.0 if retrieved[0] in set(expected) else 0.0
|
||||
|
||||
|
||||
def reciprocal_rank(retrieved: list[tuple[str, str]], expected: list[tuple[str, str]]) -> float:
|
||||
expected_set = set(expected)
|
||||
for i, page in enumerate(retrieved, start=1):
|
||||
@@ -65,8 +71,12 @@ def main() -> int:
|
||||
p = argparse.ArgumentParser()
|
||||
p.add_argument("--queries", type=Path, default=Path("eval/queries.jsonl"))
|
||||
p.add_argument("--k", type=int, default=5)
|
||||
p.add_argument("--ks", default="1,5,10,20", help="comma-separated k-curve")
|
||||
p.add_argument("--output", type=Path, default=Path("eval/results/baseline.md"))
|
||||
p.add_argument("--compare", type=Path, default=None)
|
||||
args = p.parse_args()
|
||||
ks = sorted({int(x) for x in args.ks.split(",") if x.strip()}) or [args.k]
|
||||
max_k = max(ks + [args.k])
|
||||
|
||||
if not args.queries.exists():
|
||||
print(f"queries file not found: {args.queries}")
|
||||
@@ -109,34 +119,39 @@ def main() -> int:
|
||||
rows: dict[str, dict[str, float]] = {}
|
||||
per_query: list[dict] = []
|
||||
for r in retrievers:
|
||||
mrr_sum = recall_sum = ndcg_sum = 0.0
|
||||
mrr_sum = recall_sum = ndcg_sum = p1_sum = 0.0
|
||||
elapsed_sum = 0.0
|
||||
for q in queries:
|
||||
expected = [(e["bundle_id"], e["page_id"]) for e in q["expected"]]
|
||||
t0 = time.time()
|
||||
retrieved = r.retrieve(q["query"], k=max(args.k, 10))
|
||||
retrieved = r.retrieve(q["query"], k=max(max_k, 10))
|
||||
elapsed = time.time() - t0
|
||||
mrr = reciprocal_rank(retrieved, expected)
|
||||
p1 = p_at_1(retrieved, expected)
|
||||
recall = recall_at_k(retrieved, expected, args.k)
|
||||
ndcg = ndcg_at_k(retrieved, expected, args.k)
|
||||
mrr_sum += mrr
|
||||
p1_sum += p1
|
||||
recall_sum += recall
|
||||
ndcg_sum += ndcg
|
||||
elapsed_sum += elapsed
|
||||
per_query.append({
|
||||
"retriever": r.name, "query": q["query"],
|
||||
"mrr": mrr, "recall@k": recall, "ndcg@k": ndcg,
|
||||
"mrr": mrr, "p_at_1": int(p1), "recall@k": recall, "ndcg@k": ndcg,
|
||||
"top1": list(retrieved[0]) if retrieved else None,
|
||||
"ranked": [list(p) for p in retrieved],
|
||||
"elapsed_s": round(elapsed, 3),
|
||||
})
|
||||
n = len(queries)
|
||||
rows[r.name] = {
|
||||
"P@1": p1_sum / n,
|
||||
"MRR": mrr_sum / n,
|
||||
f"Recall@{args.k}": recall_sum / n,
|
||||
f"nDCG@{args.k}": ndcg_sum / n,
|
||||
"avg_latency_s": elapsed_sum / n,
|
||||
}
|
||||
print(f" {r.name}: MRR={rows[r.name]['MRR']:.3f} "
|
||||
print(f" {r.name}: P@1={rows[r.name]['P@1']:.3f} "
|
||||
f"MRR={rows[r.name]['MRR']:.3f} "
|
||||
f"Recall@{args.k}={rows[r.name][f'Recall@{args.k}']:.3f} "
|
||||
f"nDCG@{args.k}={rows[r.name][f'nDCG@{args.k}']:.3f} "
|
||||
f"avg={rows[r.name]['avg_latency_s']*1000:.0f}ms")
|
||||
@@ -144,10 +159,10 @@ def main() -> int:
|
||||
args.output.parent.mkdir(parents=True, exist_ok=True)
|
||||
md = [f"# Retrieval eval — k={args.k}", "",
|
||||
f"_{len(queries)} hand-curated queries, generated {time.strftime('%Y-%m-%d %H:%M:%S')}_", "",
|
||||
"| Retriever | MRR | Recall@{k} | nDCG@{k} | avg latency |".replace("{k}", str(args.k)),
|
||||
"| --- | ---: | ---: | ---: | ---: |"]
|
||||
"| Retriever | P@1 | MRR | Recall@{k} | nDCG@{k} | avg latency |".replace("{k}", str(args.k)),
|
||||
"| --- | ---: | ---: | ---: | ---: | ---: |"]
|
||||
for name, m in rows.items():
|
||||
md.append(f"| `{name}` | {m['MRR']:.3f} | {m[f'Recall@{args.k}']:.3f} "
|
||||
md.append(f"| `{name}` | {m['P@1']:.3f} | {m['MRR']:.3f} | {m[f'Recall@{args.k}']:.3f} "
|
||||
f"| {m[f'nDCG@{args.k}']:.3f} | {m['avg_latency_s']*1000:.0f}ms |")
|
||||
md += ["", "## Per-query results", "",
|
||||
"| Retriever | Query | MRR | top-1 |", "| --- | --- | ---: | --- |"]
|
||||
@@ -155,7 +170,22 @@ def main() -> int:
|
||||
top1 = f"`{r['top1'][0]}/{r['top1'][1][:24]}...`" if r["top1"] else "—"
|
||||
md.append(f"| `{r['retriever']}` | {r['query'][:60]} | {r['mrr']:.3f} | {top1} |")
|
||||
args.output.write_text("\n".join(md) + "\n")
|
||||
sidecar = args.output.with_suffix(".jsonl")
|
||||
with open(sidecar, "w") as fh:
|
||||
for r in per_query:
|
||||
fh.write(json.dumps({
|
||||
"query": r["query"],
|
||||
"retriever": r["retriever"],
|
||||
"ranked": r.get("ranked") or [],
|
||||
"rr": r["mrr"],
|
||||
"p_at_1": r["p_at_1"],
|
||||
}) + "\n")
|
||||
print(f"wrote {args.output}")
|
||||
print(f"wrote {sidecar}")
|
||||
if args.compare:
|
||||
from eval.pvalue import compare, load_sidecar, render
|
||||
print()
|
||||
print(render(compare(load_sidecar(sidecar), load_sidecar(args.compare))), end="")
|
||||
return 0
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,84 @@
|
||||
"""Stdlib tests for eval metrics + the permutation test.
|
||||
|
||||
Must not open Chroma — the template has no corpus. Run with:
|
||||
|
||||
python -m unittest eval.test_metrics
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
import unittest
|
||||
|
||||
from eval.pvalue import compare, paired_permutation
|
||||
from eval.run_eval import ndcg_at_k, p_at_1, recall_at_k, reciprocal_rank
|
||||
|
||||
|
||||
A, B, C, X, Y = ("b", "a"), ("b", "b"), ("b", "c"), ("b", "x"), ("b", "y")
|
||||
|
||||
|
||||
class MetricTests(unittest.TestCase):
|
||||
def test_reciprocal_rank(self) -> None:
|
||||
# Q1: expected at rank 1
|
||||
self.assertEqual(reciprocal_rank([A, B, C], [A]), 1.0)
|
||||
# Q2: expected at rank 2
|
||||
self.assertEqual(reciprocal_rank([B, A], [A]), 0.5)
|
||||
# Q3: miss
|
||||
self.assertEqual(reciprocal_rank([X, Y], [A]), 0.0)
|
||||
|
||||
def test_p_at_1(self) -> None:
|
||||
self.assertEqual(p_at_1([A, B], [A]), 1.0)
|
||||
self.assertEqual(p_at_1([B, A], [A]), 0.0)
|
||||
self.assertEqual(p_at_1([], [A]), 0.0)
|
||||
self.assertEqual(p_at_1([A], []), 0.0)
|
||||
|
||||
def test_recall_at_k(self) -> None:
|
||||
self.assertEqual(recall_at_k([A, B, C], [A], 1), 1.0)
|
||||
self.assertEqual(recall_at_k([B, A], [A], 1), 0.0)
|
||||
self.assertEqual(recall_at_k([B, A], [A], 2), 1.0)
|
||||
self.assertEqual(recall_at_k([X, Y], [A], 5), 0.0)
|
||||
self.assertEqual(recall_at_k([A, B], [A, C], 1), 0.5)
|
||||
|
||||
def test_ndcg_at_k(self) -> None:
|
||||
self.assertEqual(ndcg_at_k([A], [A], 1), 1.0)
|
||||
# expected at rank 2: dcg = 1/log2(3), idcg = 1
|
||||
self.assertAlmostEqual(
|
||||
ndcg_at_k([B, A], [A], 2),
|
||||
(1.0 / math.log2(3)) / 1.0,
|
||||
)
|
||||
self.assertEqual(ndcg_at_k([X, Y], [A], 5), 0.0)
|
||||
|
||||
|
||||
class PermutationTests(unittest.TestCase):
|
||||
def test_identical_lists_not_significant(self) -> None:
|
||||
scores = [0.5, 1.0, 0.0, 1.0, 0.5]
|
||||
report = paired_permutation(scores, list(scores), n_resamples=2000, seed=0)
|
||||
self.assertEqual(report["Diff(A-B)"], 0.0)
|
||||
self.assertEqual(report["p_value"], 1.0)
|
||||
self.assertFalse(report["significant"])
|
||||
|
||||
def test_large_paired_difference_is_significant(self) -> None:
|
||||
a = [1.0] * 20
|
||||
b = [0.0] * 20
|
||||
report = paired_permutation(a, b, n_resamples=5000, seed=0)
|
||||
self.assertGreater(report["Diff(A-B)"], 0.9)
|
||||
self.assertLess(report["p_value"], 0.05)
|
||||
self.assertTrue(report["significant"])
|
||||
|
||||
def test_compare_joins_on_retriever_and_query(self) -> None:
|
||||
rows_a = (
|
||||
[{"retriever": "dense", "query": f"q{i}", "rr": 1.0, "p_at_1": 1} for i in range(20)]
|
||||
+ [{"retriever": "bm25", "query": "q0", "rr": 0.0, "p_at_1": 0}]
|
||||
)
|
||||
rows_b = (
|
||||
[{"retriever": "dense", "query": f"q{i}", "rr": 0.0, "p_at_1": 0} for i in range(20)]
|
||||
+ [{"retriever": "bm25", "query": "q0", "rr": 0.0, "p_at_1": 0}]
|
||||
)
|
||||
reports = compare(rows_a, rows_b, metric="rr", n_resamples=2000, seed=0)
|
||||
by_name = {r["retriever"]: r for r in reports}
|
||||
self.assertIn("dense", by_name)
|
||||
self.assertTrue(by_name["dense"]["significant"])
|
||||
self.assertFalse(by_name["bm25"]["significant"])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
+114
@@ -0,0 +1,114 @@
|
||||
"""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 = bm25_pages or dense_pages # HVM default retrieval is BM25-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())
|
||||
Reference in New Issue
Block a user