92 lines
3.2 KiB
Python
92 lines
3.2 KiB
Python
"""Bake off embedding models for the llm module (#564).
|
|
|
|
Retrieval-proxy metric: for each sampled comment, embed its chunks and query
|
|
with the comment's bib abstract (a human-written summary of the same comment).
|
|
recall@5 = fraction of comments whose own chunk ranks in the top 5 across the
|
|
pooled chunk set; plus MRR and embedding throughput. No pgvector writes — pure
|
|
in-memory numpy, so it never touches the live index.
|
|
|
|
Instruct-model selection is deliberately out of scope here: the tagging model
|
|
is chosen in P35 against the real closed vocab, not a toy prompt.
|
|
|
|
Usage:
|
|
uv run python dev/scripts/llm_bakeoff.py --docket CMS-2017-0092 --sample 120
|
|
"""
|
|
|
|
import argparse
|
|
import itertools
|
|
import time
|
|
|
|
import numpy as np
|
|
|
|
from conf.connect import bib
|
|
from llm.chunk import chunk_doc
|
|
from llm.pool import HostPool, embed_texts
|
|
from llm.source import iter_comment_docs
|
|
|
|
# (model, dims) candidates. Both serve /api/embed on the local Ollama pool.
|
|
CANDIDATES = [("nomic-embed-text", 768), ("bge-m3", 1024)]
|
|
|
|
|
|
def _normalize(mat):
|
|
return mat / np.clip(np.linalg.norm(mat, axis=1, keepdims=True), 1e-12, None)
|
|
|
|
|
|
def bakeoff(hosts, docket, sample):
|
|
store = bib()
|
|
picked = [
|
|
d
|
|
for d in itertools.islice(iter_comment_docs(store, docket=docket), sample * 3)
|
|
if d.metadata
|
|
][:sample]
|
|
# Query = each comment's abstract; skip comments without a usable abstract.
|
|
pairs = []
|
|
for d in picked:
|
|
abstract = (store.get(d.key).abstract or "").strip()
|
|
chunks = chunk_doc(d)[:4]
|
|
if abstract and chunks:
|
|
pairs.append((abstract, chunks))
|
|
print(f"# embed bake-off — docket {docket}, {len(pairs)} comments\n")
|
|
print("| model | dims | recall@5 | MRR | chunks/s |")
|
|
print("|---|---|---|---|---|")
|
|
results = {}
|
|
for model, dims in CANDIDATES:
|
|
pool = HostPool(hosts)
|
|
try:
|
|
pool.check(model)
|
|
except RuntimeError as exc:
|
|
print(f"| {model} | {dims} | SKIP | — | {exc} |")
|
|
continue
|
|
chunk_texts, owners = [], []
|
|
for i, (_, chunks) in enumerate(pairs):
|
|
chunk_texts.extend(c.text for c in chunks)
|
|
owners.extend([i] * len(chunks))
|
|
t0 = time.time()
|
|
cmat = _normalize(np.array(embed_texts(pool, model, chunk_texts)))
|
|
rate = len(chunk_texts) / (time.time() - t0)
|
|
qmat = _normalize(np.array(embed_texts(pool, model, [a for a, _ in pairs])))
|
|
owners_arr = np.array(owners)
|
|
hits = rr = 0.0
|
|
for i, q in enumerate(qmat):
|
|
order = np.argsort(cmat @ q)[::-1]
|
|
rank = int(np.where(owners_arr[order] == i)[0][0])
|
|
hits += rank < 5
|
|
rr += 1.0 / (rank + 1)
|
|
n = len(pairs)
|
|
results[model] = (hits / n, rr / n, rate)
|
|
print(f"| {model} | {dims} | {hits / n:.1%} | {rr / n:.3f} | {rate:.0f} |")
|
|
return results
|
|
|
|
|
|
def main():
|
|
ap = argparse.ArgumentParser()
|
|
ap.add_argument("--docket", default="CMS-2017-0092")
|
|
ap.add_argument("--sample", type=int, default=120)
|
|
ap.add_argument("--hosts", default="http://127.0.0.1:11434")
|
|
args = ap.parse_args()
|
|
bakeoff([h.strip() for h in args.hosts.split(",")], args.docket, args.sample)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|