This commit is contained in:
91
dev/scripts/llm_bakeoff.py
Normal file
91
dev/scripts/llm_bakeoff.py
Normal file
@@ -0,0 +1,91 @@
|
|||||||
|
"""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()
|
||||||
@@ -87,6 +87,48 @@ minus the task server it feared.
|
|||||||
off-peak on the rack card, fan out to 4090/5080 when available.
|
off-peak on the rack card, fan out to 4090/5080 when available.
|
||||||
- v1 gate: one docket end-to-end (index → query → tag → eval) before fan-out.
|
- v1 gate: one docket end-to-end (index → query → tag → eval) before fan-out.
|
||||||
|
|
||||||
|
## P33 build outcomes (2026-07-17)
|
||||||
|
|
||||||
|
P33 landed on branch `llm-p33`. Two infrastructure findings changed the plan:
|
||||||
|
|
||||||
|
**pgvector ANN index build crashes the shared Postgres (AVX-512 SIGILL).** The
|
||||||
|
mirrored `fhirworx/postgresql:latest` image (Bitnami PG18, shared by
|
||||||
|
gitea/nessie/polaris) ships a pgvector 0.8.1 `vector.so` built with AVX-512,
|
||||||
|
but the server host is a Ryzen 9 3950X (Zen 2, no AVX-512). `CREATE INDEX …
|
||||||
|
USING hnsw` **and** `ivfflat` fault with signal 4 (SIGILL) — even on 3 rows —
|
||||||
|
and the fault kills every backend, dropping the whole server (and its other
|
||||||
|
databases) into recovery. Index-free **exact search works fine** (query-time
|
||||||
|
distance dispatch is correct; only the bulk index-build path faults). Decision:
|
||||||
|
ANN index creation is gated behind `[llm].build_ann_index` (**default false**);
|
||||||
|
the module ships on exact search. Restoring HNSW needs the image's pgvector
|
||||||
|
rebuilt without AVX-512 (thin image `FROM fhirworx/postgresql:latest` replacing
|
||||||
|
`vector.so` with a `-march=x86-64-v2`/`znver2` build) — filed as a P33
|
||||||
|
follow-up issue. Exact search on the pilot (15,852 vectors) answers in ~1.2 s
|
||||||
|
for two queries, so this is not urgent at pilot scale.
|
||||||
|
|
||||||
|
**Embed model bake-off (#564).** Retrieval-proxy metric (comment abstract as
|
||||||
|
query against its own pooled chunks, 120 pilot comments):
|
||||||
|
|
||||||
|
| model | dims | recall@5 | MRR | chunks/s (1× 3060) |
|
||||||
|
|---|---|---|---|---|
|
||||||
|
| nomic-embed-text | 768 | 86.7% | 0.850 | 105 |
|
||||||
|
| bge-m3 | 1024 | 93.3% | 0.893 | 18 |
|
||||||
|
|
||||||
|
bge-m3 retrieves better (+6.6 pp recall@5) but is ~6× slower — ~25 h vs ~4 h to
|
||||||
|
embed the full 164k-comment corpus on one card. **P33 default stays
|
||||||
|
`nomic-embed-text` (768-dim)**: the pilot is validated on it, 86.7% recall is
|
||||||
|
solid, and it keeps the delivered index internally consistent. **Recommended
|
||||||
|
upgrade for the corpus-wide run: bge-m3**, adopted at P35/scale-up time when the
|
||||||
|
4090/5080 pool absorbs the throughput cost — switching means `embed_dim = 1024`
|
||||||
|
and a `--force` re-index. Instruct-model selection is deferred to P35, where it
|
||||||
|
is judged against the real closed vocab rather than a toy prompt.
|
||||||
|
|
||||||
|
**Pilot result:** docket CMS-2017-0092 indexed end-to-end — 1,618 comments →
|
||||||
|
15,852 chunks; a killed run resumed with 25 already-done items skipped
|
||||||
|
(resumability proven); semantic queries return on-topic comments (telehealth
|
||||||
|
originating-site §1834(m) exemptions; documentation-burden reduction). The v1
|
||||||
|
gate (one docket, index → query) is met; tag → eval arrives in P35.
|
||||||
|
|
||||||
## Relationship to P30
|
## Relationship to P30
|
||||||
|
|
||||||
#254 (position classification + thematic tagging) gets *implemented on* this
|
#254 (position classification + thematic tagging) gets *implemented on* this
|
||||||
|
|||||||
Reference in New Issue
Block a user