"""Vidul's fine-tuned dense retrieval for medical-code search. The runtime uses the original BGE_FT_VA model and FAISS indexes with no score blend or fallback model. Indexes are loaded lazily by category so a CPU Space does not pay the memory cost for categories nobody has queried. """ from __future__ import annotations import bisect import csv import hashlib 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: ") 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]] = {} 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) -> 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")) digest = hashlib.sha256(meta_path.read_bytes()).hexdigest() 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")) with (directory / "meta.csv").open(newline="", encoding="utf-8") as handle: records = list(csv.DictReader(handle)) 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))}) self._loaded[category] = loaded self._suggesters[category] = SpellSuggester( record["description"] for source in loaded for record in source["records"]) return loaded 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) -> dict: sources = self._load_category(category) query_vector = self.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": "BGE_FT_VA", "suggestion": self._suggesters[category].suggest(query.strip()), "low_confidence": (not results or results[0]["relevance"] < LOW_CONFIDENCE.get(category, 0.75)), "results": results, }