feat(llm): embed bake-off harness + P33 outcomes in spec (refs #564 #567)

This commit is contained in:
kert
2026-07-17 14:37:34 -04:00
parent b2c129fd06
commit c85d384f32
2 changed files with 133 additions and 0 deletions

View 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()

View File

@@ -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