"""Dense retrieval for medical-code search. The default model is Vidul's BGE_FT_VA with its original FAISS indexes, with no score blend or fallback model. Comparison models attach their own indexes (`attach`): one index per searched index name, built over exactly the default rows in the same order, so every model shares one copy of the row metadata and of every per-row sidecar. Each attached index is checked against the default meta.csv fingerprint before it is trusted. Indexes load per category, but the app warms every category at startup (`warm`) rather than on first use: deferring the read only moved a multi-hundred-megabyte read of a lazily fetched network mount onto a visitor's request, which is the read pattern that mount fails on. The saving was never real either, since every category is queried in an ordinary session. """ from __future__ import annotations import bisect import csv import hashlib import io import json import os from collections import Counter import faiss import numpy as np from . import paths from . import provenance from .mappings import MappingStore, get_store from .retriever import DenseEmbedder from .spelling import SpellSuggester # Top-score floor below which a result set is flagged low-confidence, per # category — the scales differ (nonsense probes reach ~0.87 cosine against the # 430k-row lab index but only ~0.76 against CPT). Calibrated 2026-08-05 against # demo-bundle real queries vs nonsense probes: every observed real query clears # its floor with margin; refine from gold-pair score distributions as they # accumulate (scripts/eval_gold_pairs.py). def _low_conf(category: str, default: float) -> float: return float(os.environ.get(f"ENCODE_LOW_CONFIDENCE_{category.upper()}", default)) LOW_CONFIDENCE = {"diagnosis": _low_conf("diagnosis", 0.75), "medication": _low_conf("medication", 0.74), "lab": _low_conf("lab", 0.80), "procedure": _low_conf("procedure", 0.78)} CATEGORY_INDEXES = { "diagnosis": [("icd_index", None)], # The NDC index is not searched directly: NDC codes reach results through # mappings on local drug codes, not as a peer code system. "medication": [("med_index", "Local Drug SID")], "lab": [("labchem_index", "LabChemTestSID")], "procedure": [("cpt_index", "CPT")], } # The prebuilt indexes key their rows by description, so an ICD-9 and an # ICD-10 code sharing a description land in one row carrying both codes and # both labels ("039.2,A42.1" / "ICD10 | ICD9"). That label describes the row, # not either code, and no code is ever both versions. These sidecars give each # code in such a row its own version so results carry one code system each and # the merged label never reaches a caller. Built by # scripts/build_icd_index_versions.py; see it for how the split is resolved. ROW_VERSION_SIDECARS = {"icd_index": paths.ICD_INDEX_VERSIONS} def _normalize_code(value: str) -> str: """One spelling for a code, so N18.9, n189 and " N18.9 " are one key.""" return str(value).strip().upper().replace(".", "").replace(" ", "") def make_embedder() -> DenseEmbedder: """Reproduce Vidul's SentenceTransformer query convention on CPU.""" if not paths.MODEL_DIR.exists(): raise RuntimeError(f"Missing Vidul fine-tuned model: {paths.MODEL_DIR}") return DenseEmbedder(str(paths.MODEL_DIR), pooling="mean", query_prefix="passage: ") # The default indexes predate per-model configs: they name their model this way # and leave the query convention to make_embedder(). DEFAULT_INDEX_MODEL = "BGE_FT_VA" # Per-category low-confidence floors for an attached model, written beside its # indexes by scripts/calibrate_low_confidence.py. A model without the file # flags nothing: a floor calibrated for another model's score scale would be # wrong in a way no one could see. LOW_CONFIDENCE_FILE = "low_confidence.json" def index_config(directory) -> dict: """The config.json beside one attached index.""" return json.loads((directory / "config.json").read_text()) class CodeSearchEngine: # Adding a mapping to a category = adding a stamp method and one entry here. STAMPERS = {"lab": "_stamp_lab", "diagnosis": "_stamp_diagnosis", "procedure": "_stamp_procedure"} def __init__(self, embedder: DenseEmbedder, store: MappingStore | None = None): self.embedder = embedder self.store = store or get_store() self._loaded: dict[str, list[dict]] = {} self._suggesters: dict[str, SpellSuggester] = {} self._systems: dict[str, list[dict]] = {} self._rbcs_sizes: dict[str, int] | None = None self._code_order: dict[str, tuple[list[str], dict]] = {} # meta.csv fingerprint per default index, which attached indexes must match. self._meta_sha: dict[str, str] = {} # Attached models: id -> {"dir", "embedder", "floors"}, and their indexes # per category, aligned with the default sources. self._models: dict[str, dict] = {} self._model_indexes: dict[tuple[str, str], list] = {} missing = [name for specs in CATEGORY_INDEXES.values() for name, _ in specs if not self._complete_index(paths.CODE_INDEX_DIR / name)] if missing: raise RuntimeError( f"Missing Vidul FAISS index assets under {paths.CODE_INDEX_DIR}: " f"{', '.join(missing)}") @staticmethod def _complete_index(directory) -> bool: return all((directory / name).exists() for name in ("index.faiss", "meta.csv", "config.json")) @staticmethod def _load_row_versions(directory_name: str, meta_path, rows: int, digest: str) -> dict[int, list[tuple[str, str]]]: """Per-code ICD versions for the rows that pack two code systems. Keyed by row position, which is only meaningful against the exact meta.csv the sidecar was built from, so the fingerprint is checked rather than trusted: a regenerated index with a stale sidecar would otherwise silently relabel unrelated codes. """ path = ROW_VERSION_SIDECARS.get(directory_name) if path is None: return {} if not path.exists(): raise RuntimeError( f"Missing per-code version sidecar {path}; build it with " f"scripts/build_icd_index_versions.py") payload = json.loads(path.read_text(encoding="utf-8")) if payload.get("meta_rows") != rows or payload.get("meta_sha256") != digest: raise RuntimeError( f"{path} was built from a different {meta_path}; rerun " f"scripts/build_icd_index_versions.py") return {int(position): [(code, version) for code, version in pairs] for position, pairs in payload["rows"].items()} @staticmethod def _searchable(source: dict) -> int: """Results the source can return, which is its row count only while no row splits. A split row yields one result per code, so counting rows would report a smaller corpus than the code-system facets add up to.""" split = source["row_versions"] return (source["index"].ntotal - len(split) + sum(len(pairs) for pairs in split.values())) @staticmethod def _variants(source: dict, position: int, record: dict) -> list[tuple[str, str]]: """The (code, code_type) pairs one index row contributes to results. An ordinary row contributes itself. A row that packs an ICD-9 and an ICD-10 code under one description contributes one pair per code, each carrying its own resolved version, so a caller never sees a single code labelled with two systems. Splitting only relabels: the codes a category serves are exactly the codes in the source index. """ split = source["row_versions"].get(position) if split: return split return [(record["code"], source["code_type"] or record.get("version") or "Diagnosis")] def categories(self) -> list[str]: return list(CATEGORY_INDEXES) def _load_category(self, category: str) -> list[dict]: if category not in CATEGORY_INDEXES: raise KeyError(category) if category in self._loaded: return self._loaded[category] loaded = [] for directory_name, fixed_code_type in CATEGORY_INDEXES[category]: directory = paths.CODE_INDEX_DIR / directory_name config = json.loads((directory / "config.json").read_text()) if config.get("model") != "BGE_FT_VA" or config.get("dim") != 1024: raise RuntimeError(f"Unexpected retrieval config: {directory / 'config.json'}") index = faiss.read_index(str(directory / "index.faiss")) # Read once: the bytes give both the rows and the fingerprint that # the row-version sidecar and every attached index are checked by. meta_bytes = (directory / "meta.csv").read_bytes() digest = hashlib.sha256(meta_bytes).hexdigest() self._meta_sha[directory_name] = digest records = list(csv.DictReader(io.StringIO(meta_bytes.decode("utf-8"), newline=""))) if index.ntotal != len(records): raise RuntimeError( f"Index/metadata mismatch in {directory}: {index.ntotal} != {len(records)}") loaded.append({"index": index, "records": records, "code_type": fixed_code_type, "row_versions": self._load_row_versions( directory_name, directory / "meta.csv", len(records), digest)}) self._loaded[category] = loaded self._suggesters[category] = SpellSuggester( record["description"] for source in loaded for record in source["records"]) return loaded def warm(self, category: str, model_id: str | None = None) -> None: """Force this category's index in, so no request ever loads one. Loading on first use put a multi-hundred-megabyte read of a lazily fetched network mount on a visitor's request, concurrently with whatever else was reading it. That is the read pattern the mount fails on, so every category is pulled in at startup instead, one at a time. """ self._load_category(category) if model_id is not None: self._load_model_category(model_id, category) # -- attached comparison models -------------------------------------------- def attach(self, model_id: str, index_dir, embedder: DenseEmbedder) -> None: """Register a comparison model's indexes; `warm` then loads them.""" missing = [name for specs in CATEGORY_INDEXES.values() for name, _ in specs if not self._complete_index(index_dir / name)] if missing: raise RuntimeError( f"Missing FAISS index assets under {index_dir}: {', '.join(missing)}") floors_path = index_dir / LOW_CONFIDENCE_FILE floors = json.loads(floors_path.read_text())["floors"] if floors_path.exists() else {} self._models[model_id] = {"dir": index_dir, "embedder": embedder, "floors": floors} def models(self) -> list[str]: return list(self._models) def _load_model_category(self, model_id: str, category: str) -> list: key = (model_id, category) if key in self._model_indexes: return self._model_indexes[key] model = self._models.get(model_id) if model is None: raise LookupError(f"Model '{model_id}' is not loaded for code search.") sources = self._load_category(category) indexes = [] for (directory_name, _), source in zip(CATEGORY_INDEXES[category], sources): directory = model["dir"] / directory_name config = index_config(directory) if config.get("model_id") != model_id: raise RuntimeError(f"{directory} was built for '{config.get('model_id')}', " f"not '{model_id}'") if config.get("meta_sha256") != self._meta_sha[directory_name]: raise RuntimeError(f"{directory} was built from a different meta.csv than " f"{paths.CODE_INDEX_DIR / directory_name}") index = faiss.read_index(str(directory / "index.faiss")) if index.ntotal != source["index"].ntotal: raise RuntimeError(f"Index/metadata mismatch in {directory}: " f"{index.ntotal} != {source['index'].ntotal}") indexes.append(index) self._model_indexes[key] = indexes return indexes def _sources_for(self, category: str, model_id: str | None) -> list[dict]: """The default sources, with the model's index in place of the default one.""" sources = self._load_category(category) if model_id is None: return sources indexes = self._load_model_category(model_id, category) return [{**source, "index": index} for source, index in zip(sources, indexes)] def records(self, category: str) -> list[dict]: """Metadata rows in the shape consumed by the optional graph view.""" rows = [] for source in self._load_category(category): for position, record in enumerate(source["records"]): for code, code_type in self._variants(source, position, record): rows.append({"code": code, "description": record["description"], "code_type": code_type}) return rows # -- per-category mapping stamps (adding a mapping = adding a function) -- # A lab index row packs every LabChemTestSID that shares one description, # and 73 of them is ordinary. Requiring all of them to name the same LOINC # threw away the row whenever a single SID dissented, so "hemoglobin a1c", # where 66 of 67 mapped SIDs say 4548-4, showed no mapping at all. A clear # majority now carries the row and is reported as derived, since the row's # assertion is then ENCODE's and not the crosswalk's. AGREEMENT = 2 / 3 def _stamp_lab(self, result: dict, record: dict) -> None: # Each entry is paired with the SID that carries it: the row is cited # and graphed by a code that actually holds the mapping, which is not # always the first one in the field. entries = [(code.strip(), entry) for code, entry in ((code, self.store.loinc_of(code)) for code in str(record["code"]).split(",")) if entry] result["graphable"] = "LOINC" in result["code_type"].upper() or bool(entries) if not entries: return counts = Counter(entry["loinc"] for _, entry in entries) loinc, top = counts.most_common(1)[0] split = len(counts) > 1 if split and top / len(entries) < self.AGREEMENT: # Genuinely mixed: the codes in this row are not one test. Say so # rather than asserting one of them. result["graphable"] = "LOINC" in result["code_type"].upper() result["mapping_conflict"] = {"targets": len(counts), "codes": len(entries)} return agreeing = [(code, entry) for code, entry in entries if entry["loinc"] == loinc] # Cite the strongest claim on that LOINC: a code the crosswalk mapped # outright, if the row has one, and only otherwise a derived code. The # row still counts as derived when the codes disagreed, because then # the choice between them is ENCODE's rather than the crosswalk's. code, entry = next((pair for pair in agreeing if not pair[1].get("derived")), agreeing[0]) derived = bool(entry.get("derived")) or split result["mapped_loinc"] = loinc result["mapped_code"] = code if derived: result["mapping_derived"] = True result["mapping_provenance"] = provenance.line( entry.get("source"), derived=derived, match="merged_majority" if split else entry.get("match"), target=entry.get("loinc_version"), detail=f"{top} of {len(entries)}" if split else None) def _stamp_diagnosis(self, result: dict, record: dict) -> None: # Read off the result, not the row: a row holding both ICD versions has # already been split into one result per code, and each must be stamped # with the phecode for its own code under its own version. Stamping from # the row would hand an ICD-10 result its sibling's ICD-9 assignment. versions = [part.strip() for part in result["code_type"].split("|")] # A row can still pack several codes of one version ("585.9,585.6"). # The phecode maps are keyed on a single dotted code, so the whole # field never matched and 4,394 packed rows carried no phecode at all. # Each code is tried in turn. for code in str(result["code"]).split(","): phecodes = self.store.phecodes_for(code, versions) if phecodes: result["phecodes"] = {**phecodes, "provenance": provenance.phecode_line(phecodes)} result.setdefault("mapped_code", code.strip()) return def rbcs_family_sizes(self) -> dict[str, int]: """Codes per RBCS family, counted over the procedure codes served here. CMS's mapping file and the CPT index are not the same set. Counting the file gave "Arthroscopy - Lower Extremity, 46 codes" on the detail page while the graph, drawn from the index, said "Showing 26 of 45": two numbers for one family, which is the confusion the family size was added to remove. This counts the set the graph draws, by the same rule (validated membership, each code once), so the two agree by construction rather than by luck. """ if self._rbcs_sizes is None: sizes: dict[str, int] = {} seen: set[str] = set() for row in self.records("procedure"): if "CPT" not in str(row["code_type"]): continue for code in str(row["code"]).split(","): code = code.strip() if not code or code in seen: continue seen.add(code) entry = self.store.rbcs_validated_for(code) or {} group = entry.get("family") or entry.get("subcategory") if group: sizes[group] = sizes.get(group, 0) + 1 self._rbcs_sizes = sizes return self._rbcs_sizes def _stamp_procedure(self, result: dict, record: dict) -> None: # Same merged-row rule as the other two: the first code in the row is # not always the one CMS assigned a group to. rbcs = code = None for candidate in str(record["code"]).split(","): rbcs = self.store.rbcs_for(candidate) if rbcs: code = candidate.strip() break if rbcs: result["mapped_code"] = code result["rbcs"] = {k: rbcs[k] for k in ("category", "subcategory", "family", "major")} # How broad the family is travels with the group: a name shared by # two codes is normal, and the size is what stops it reading as a # mistake. Absent when this code reached its group by a derived # mapping, because then the graph has no validated family to draw # and there is no number both surfaces could agree on. group = rbcs.get("family") or rbcs.get("subcategory") size = self.rbcs_family_sizes().get(group) if group else None if size: result["rbcs"]["family_size"] = size result["graphable"] = True if rbcs.get("derived"): result["mapping_derived"] = True result["mapping_provenance"] = provenance.line( rbcs.get("source"), derived=bool(rbcs.get("derived")), match=rbcs.get("match")) def systems(self, category: str) -> list[dict]: """Distinct code systems in a category with row counts, for the pre-search restriction checkboxes. Computed once per category.""" cached = self._systems.get(category) if cached is None: counts: dict[str, int] = {} for source in self._load_category(category): if source["code_type"]: counts[source["code_type"]] = (counts.get(source["code_type"], 0) + len(source["records"])) else: # Counted per code rather than per row: a row that packs an # ICD-9 and an ICD-10 code contributes one result to each # system, which is exactly what checking that box returns. for position, record in enumerate(source["records"]): for _, ct in self._variants(source, position, record): counts[ct] = counts.get(ct, 0) + 1 cached = [{"code_type": t, "count": n} for t, n in sorted(counts.items(), key=lambda kv: -kv[1])] self._systems[category] = cached return cached def _codes_in_order(self, category: str): """Every distinct normalized code, sorted, with the rows that carry it. This is what makes the lookup a lookup: the neighbours of a code are the codes beside it in its own numbering, so 1000 sits next to 1001 and 10001, the way it does in the code book. Built once per category on first use and held; the largest category packs 640,871 codes. """ cached = self._code_order.get(category) if cached: return cached by_code: dict[str, list[tuple[int, int, str, str]]] = {} for source_index, source in enumerate(self._load_category(category)): for position, record in enumerate(source["records"]): for value, code_type in self._variants(source, position, record): for part in str(value).split(","): norm = _normalize_code(part) if norm: by_code.setdefault(norm, []).append( (source_index, position, value, code_type)) cached = (sorted(by_code), by_code) self._code_order[category] = cached return cached def lookup(self, category: str, code: str, k: int = 50) -> dict: """The code's own rows, then the codes beside it in code order. Search embeds the query as text and scores it against descriptions, so a code never retrieves itself: "90999" ranks against the words in the index and returns "time", "introduction", "home". Anyone arriving with a code in hand is asking a different question, and it is answered from the code column. Matching ignores case, surrounding space, and dots, so a copied " n189 " finds N18.9. Related rows are the sorted neighbourhood of the query, nearest first, where nearest means the longer shared prefix: for 1000 that is 10001 before 1001 before 1100. A partial or unknown code therefore still lands in the right part of the code book and shows what is there. """ sources = self._load_category(category) stamp_name = self.STAMPERS.get(category) stamp = getattr(self, stamp_name) if stamp_name else None ordered, by_code = self._codes_in_order(category) wanted = _normalize_code(code) seen: set[tuple[int, int, str, str]] = set() def rows_for(norm: str, label: str) -> list[dict]: out = [] for entry in by_code.get(norm, []): if entry in seen: # a packed row answers for one code once continue seen.add(entry) source_index, position, value, code_type = entry record = sources[source_index]["records"][position] # No relevance: nothing here was ranked, and printing a score # would invent one. row = {"code": value, "code_type": code_type, "description": record["description"], "relevance": None, "match": label} if stamp: stamp(row, record) out.append(row) return out results = rows_for(wanted, "exact")[:k] if wanted else [] # Walk outward from the query's place in the sorted code list, # taking the side that shares the longer prefix with it first. `k` is # the same result-count setting the search obeys and caps the whole # page, so nearby rows fill whatever the exact matches left of it. related_cap = max(0, k - len(results)) related: list[dict] = [] if wanted: def shared(norm: str) -> int: n = 0 for a, b in zip(norm, wanted): if a != b: break n += 1 return n left = bisect.bisect_left(ordered, wanted) - 1 right = bisect.bisect_right(ordered, wanted) while len(related) < related_cap and (left >= 0 or right < len(ordered)): if right >= len(ordered): take_right = False elif left < 0: take_right = True else: take_right = shared(ordered[right]) >= shared(ordered[left]) norm = ordered[right if take_right else left] if take_right: right += 1 else: left -= 1 related.extend(rows_for(norm, "nearby code")) related = related[:related_cap] rows = results + related for rank, row in enumerate(rows, 1): row["rank"] = rank return { "query": code.strip(), "category": category, # Truthful about what was found: a lookup that hit the code says # exact, one that only found neighbours says nearest. "match": "exact" if results else "nearest", "total": sum(self._searchable(source) for source in sources), "count": len(rows), "exact_count": len(results), "related_count": len(related), "results": rows, } def search(self, category: str, query: str, k: int = 50, systems: set[str] | None = None, model_id: str | None = None) -> dict: """Rank codes for `query`. `model_id` names an attached model; None is the default model.""" sources = self._sources_for(category, model_id) if model_id is None: embedder, floor = self.embedder, LOW_CONFIDENCE.get(category, 0.75) else: model = self._models[model_id] embedder, floor = model["embedder"], model["floors"].get(category) query_vector = embedder.encode([query.strip()]).astype(np.float32) stamp_name = self.STAMPERS.get(category) stamp = getattr(self, stamp_name) if stamp_name else None # A restricted system may be a thin slice of a source, so over-fetch # before filtering; FAISS neighbor count barely affects a flat scan. fetch_k = k if not systems else max(k * 10, 500) candidates = [] for source in sources: scores, indices = source["index"].search( query_vector, min(fetch_k, source["index"].ntotal)) for score, index_position in zip(scores[0], indices[0]): if index_position < 0: continue record = source["records"][int(index_position)] # One row can carry more than one code system; each is its own # result, so every code_type below names exactly one system. for code, code_type in self._variants(source, int(index_position), record): if systems and code_type not in systems: continue result = { "code": code, "code_type": code_type, "description": record["description"], "_score": float(np.clip(score, -1.0, 1.0)), } if stamp: stamp(result, record) candidates.append(result) candidates.sort(key=lambda row: row["_score"], reverse=True) results = candidates[:k] for rank, result in enumerate(results, 1): result["relevance"] = round(result.pop("_score"), 4) result["rank"] = rank return { "query": query, "category": category, "total": sum(self._searchable(source) for source in sources), "count": len(results), "model": model_id or DEFAULT_INDEX_MODEL, "suggestion": self._suggesters[category].suggest(query.strip()), "low_confidence": (not results or (floor is not None and results[0]["relevance"] < floor)), "results": results, }