Spaces:
Running
Running
Claude Opus 5.5
Serve four embedding models: BGE VA-tuned (default), BGE base, Harrier PARSE-tuned, Harrier base
5e5072d verified | """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)}") | |
| def _complete_index(directory) -> bool: | |
| return all((directory / name).exists() | |
| for name in ("index.faiss", "meta.csv", "config.json")) | |
| 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()} | |
| 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())) | |
| 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, | |
| } | |