Spaces:

anurag-raapid
/

query-entity-export

Runtime error

query-entity-export / src /data /query_entity_linker.py

Upload folder using huggingface_hub

fd1eb5a verified 4 months ago

15.8 kB

"""Link clinical evidence text to UMLS graph nodes.
Primary backend: surface-form dictionary built from the filtered KG (``id_maps.parquet``)
plus optional enrichment from ``mrxns_eng`` / ``mrsty`` in the UMLS DuckDB.
Optional backend: QuickUMLS (when ``leveldb`` + ``quickumls`` are installed).
"""

from __future__ import annotations

import json
import re
import unicodedata
from pathlib import Path

import numpy as np
import polars as pl

from config.paths import (
    CUI_TO_GRAPH_IDX_JSON,
    GRAPH_ID_MAPS_PARQUET,
    QUERY_LINKER_CACHE_JSON,
    UMLS_DUCKDB_PATH,
)

# Clinical semantic types (diseases, drugs, symptoms, procedures, anatomy, findings).
CLINICAL_SEMTYPES: frozenset[str] = frozenset(
    {
        "T047",  # Disease or Syndrome
        "T048",  # Mental or Behavioral Dysfunction
        "T184",  # Sign or Symptom
        "T121",  # Pharmacologic Substance
        "T109",  # Organic Chemical
        "T123",  # Biologically Active Substance
        "T195",  # Antibiotic
        "T200",  # Clinical Drug
        "T023",  # Body Part, Organ, or Organ Component
        "T031",  # Body Substance
        "T033",  # Finding
        "T034",  # Laboratory or Test Result
        "T060",  # Diagnostic Procedure
        "T061",  # Therapeutic or Preventive Procedure
        "T037",  # Injury or Poisoning
        "T046",  # Pathologic Function
        "T191",  # Neoplastic Process
    }
)

PREFERRED_SABS: frozenset[str] = frozenset(
    {
        "SNOMEDCT_US",
        "MSH",
        "RXNORM",
        "MEDCIN",
        "NCI",
        "LNC",
        "ICD10CM",
        "ICD10",
        "ICD9CM",
        "NDDF",
        "MMSL",
        "VANDF",
    }
)

STOPWORDS: frozenset[str] = frozenset(
    {
        "a",
        "an",
        "the",
        "and",
        "or",
        "of",
        "in",
        "on",
        "at",
        "to",
        "for",
        "with",
        "without",
        "by",
        "from",
        "is",
        "are",
        "was",
        "were",
        "be",
        "been",
        "being",
        "has",
        "have",
        "had",
        "not",
        "no",
        "due",
        "associated",
        "secondary",
        "primary",
        "unspecified",
        "other",
        "specified",
        "history",
        "personal",
        "patient",
        "given",
        "then",
        "there",
        "will",
        "some",
        "any",
        "all",
        "both",
        "each",
        "than",
        "that",
        "this",
        "these",
        "those",
        "into",
        "during",
        "after",
        "before",
        "while",
        "when",
        "where",
        "which",
        "who",
        "whom",
        "whose",
        "if",
        "as",
        "but",
        "so",
        "such",
        "via",
        "per",
        "also",
        "only",
        "just",
        "very",
        "more",
        "most",
        "less",
        "least",
        "new",
        "old",
        "acute",
        "chronic",
        "stage",
        "type",
        "status",
        "present",
        "absent",
        "possible",
        "probable",
        "suspected",
    }
)

_SURFACE_RE = re.compile(r"[^a-z0-9]+")


def normalize_surface(text: str) -> str:
    """Lowercase, strip accents, collapse punctuation to spaces."""
    text = unicodedata.normalize("NFKD", text)
    text = "".join(ch for ch in text if not unicodedata.combining(ch))
    text = text.lower().strip()
    text = _SURFACE_RE.sub(" ", text)
    return " ".join(text.split())


def _tokenize(text: str) -> list[str]:
    return [t for t in normalize_surface(text).split() if t and t not in STOPWORDS]


def build_cui_to_graph_idx(
    id_maps_parquet: Path = GRAPH_ID_MAPS_PARQUET,
    max_graph_idx: int | None = None,
) -> dict[str, int]:
    """Map CUI -> best graph row index (highest degree node per CUI)."""
    scan = pl.scan_parquet(id_maps_parquet).select("idx", "cui", "degree")
    if max_graph_idx is not None:
        scan = scan.filter(pl.col("idx") < max_graph_idx)
    df = (
        scan.filter(pl.col("cui").is_not_null())
        .sort(["cui", "degree"], descending=[False, True])
        .group_by("cui")
        .agg(pl.col("idx").first())
        .collect()
    )
    return dict(zip(df["cui"].to_list(), df["idx"].to_list()))


def build_clinical_cuis(
    cui_to_idx: dict[str, int],
    umls_duckdb: Path | None = UMLS_DUCKDB_PATH,
) -> set[str]:
    """Filter CUIs to clinical semantic types when DuckDB is available."""
    if umls_duckdb is None or not Path(umls_duckdb).exists():
        return set(cui_to_idx)
    import duckdb

    con = duckdb.connect(str(umls_duckdb), read_only=True)
    con.register("graph_cuis", pl.DataFrame({"cui": list(cui_to_idx)}))
    sem = ",".join(f"'{s}'" for s in CLINICAL_SEMTYPES)
    rows = con.execute(
        f"""
        SELECT DISTINCT g.cui
        FROM graph_cuis g
        JOIN mrsty s ON g.cui = s.CUI
        WHERE s.STY IN ({sem})
        """
    ).fetchall()
    con.close()
    return {r[0] for r in rows} if rows else set(cui_to_idx)


def build_surface_lookup(
    cui_to_idx: dict[str, int],
    clinical_cuis: set[str],
    id_maps_parquet: Path = GRAPH_ID_MAPS_PARQUET,
    umls_duckdb: Path | None = UMLS_DUCKDB_PATH,
    max_graph_idx: int | None = None,
    min_surface_len: int = 3,
) -> dict[str, list[tuple[str, float]]]:
    """Build normalized surface form -> [(cui, score)] from KG + optional MRXNS."""
    lookup: dict[str, dict[str, float]] = {}

    def add(surface: str, cui: str, score: float) -> None:
        if cui not in clinical_cuis:
            return
        norm = normalize_surface(surface)
        if len(norm) < min_surface_len:
            return
        bucket = lookup.setdefault(norm, {})
        bucket[cui] = max(bucket.get(cui, 0.0), score)

    scan = pl.scan_parquet(id_maps_parquet).select("idx", "cui", "sab", "str")
    if max_graph_idx is not None:
        scan = scan.filter(pl.col("idx") < max_graph_idx)
    df = scan.filter(
        pl.col("cui").is_not_null()
        & pl.col("str").is_not_null()
        & (pl.col("str").str.len_chars() >= min_surface_len)
    ).collect()

    for cui, sab, surface in zip(
        df["cui"].to_list(), df["sab"].to_list(), df["str"].to_list()
    ):
        score = 1.0 if sab in PREFERRED_SABS else 0.85
        add(surface, cui, score)

    if umls_duckdb is not None and Path(umls_duckdb).exists():
        import duckdb

        con = duckdb.connect(str(umls_duckdb), read_only=True)
        con.register("graph_cuis", pl.DataFrame({"cui": list(clinical_cuis)}))
        sem = ",".join(f"'{s}'" for s in CLINICAL_SEMTYPES)
        rows = con.execute(
            f"""
            SELECT DISTINCT m.CUI, m.NSTR
            FROM mrxns_eng m
            JOIN graph_cuis g ON m.CUI = g.cui
            JOIN mrsty s ON m.CUI = s.CUI
            WHERE s.STY IN ({sem})
              AND length(m.NSTR) >= {int(min_surface_len)}
            """
        ).fetchall()
        con.close()
        for cui, nstr in rows:
            surface = " ".join(nstr.split())
            add(surface, cui, 0.9)

    return {k: sorted(v.items(), key=lambda x: -x[1]) for k, v in lookup.items()}


def save_cui_to_graph_idx(path: Path, cui_to_idx: dict[str, int]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(cui_to_idx), encoding="utf-8")


def load_cui_to_graph_idx(path: Path = CUI_TO_GRAPH_IDX_JSON) -> dict[str, int]:
    payload = json.loads(path.read_text(encoding="utf-8"))
    return {str(k): int(v) for k, v in payload.items()}


def ensure_cui_to_graph_idx(
    path: Path = CUI_TO_GRAPH_IDX_JSON,
    id_maps_parquet: Path = GRAPH_ID_MAPS_PARQUET,
    rebuild: bool = False,
) -> dict[str, int]:
    """Build or load CUI → graph row index (full graph, highest-degree node per CUI)."""
    if path.exists() and not rebuild:
        return load_cui_to_graph_idx(path)
    cui_to_idx = build_cui_to_graph_idx(id_maps_parquet, max_graph_idx=None)
    save_cui_to_graph_idx(path, cui_to_idx)
    return cui_to_idx


def save_linker_cache(
    path: Path,
    cui_to_idx: dict[str, int],
    surface_lookup: dict[str, list[tuple[str, float]]],
) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    payload = {
        "cui_to_idx": cui_to_idx,
        "surface_lookup": {k: list(v) for k, v in surface_lookup.items()},
    }
    path.write_text(json.dumps(payload), encoding="utf-8")


def load_linker_cache(
    path: Path,
) -> tuple[dict[str, int], dict[str, list[tuple[str, float]]]]:
    payload = json.loads(path.read_text(encoding="utf-8"))
    surface_lookup = {
        k: [(cui, float(score)) for cui, score in v]
        for k, v in payload["surface_lookup"].items()
    }
    return payload["cui_to_idx"], surface_lookup


class QueryEntityLinker:
    """Map query text to graph node indices with confidence scores."""

    def __init__(
        self,
        *,
        id_maps_parquet: Path = GRAPH_ID_MAPS_PARQUET,
        umls_duckdb: Path | None = UMLS_DUCKDB_PATH,
        cache_path: Path = QUERY_LINKER_CACHE_JSON,
        quickumls_path: Path | None = None,
        max_graph_idx: int | None = None,
        rebuild_cache: bool = False,
        min_confidence: float = 0.7,
        max_ngram: int = 6,
        quickumls_only: bool = False,
        skip_graph_cui_filter: bool = False,
    ):
        self.min_confidence = min_confidence
        self.max_ngram = max_ngram
        self.quickumls_only = quickumls_only
        self.skip_graph_cui_filter = skip_graph_cui_filter
        self._quickumls = None

        if quickumls_only and skip_graph_cui_filter:
            self.cui_to_idx = {}
            self.surface_lookup = {}
        elif quickumls_only:
            self.cui_to_idx = ensure_cui_to_graph_idx(
                rebuild=rebuild_cache,
                id_maps_parquet=id_maps_parquet,
            )
            self.surface_lookup = {}
        elif cache_path.exists() and not rebuild_cache:
            self.cui_to_idx, self.surface_lookup = load_linker_cache(cache_path)
        else:
            # Build full-graph linker cache; OOB embedding rows are filtered at pool time.
            cui_to_idx = build_cui_to_graph_idx(id_maps_parquet, max_graph_idx=None)
            clinical_cuis = build_clinical_cuis(cui_to_idx, umls_duckdb)
            self.surface_lookup = build_surface_lookup(
                cui_to_idx,
                clinical_cuis,
                id_maps_parquet=id_maps_parquet,
                umls_duckdb=umls_duckdb,
                max_graph_idx=None,
            )
            self.cui_to_idx = {
                c: i for c, i in cui_to_idx.items() if c in clinical_cuis
            }
            save_linker_cache(cache_path, self.cui_to_idx, self.surface_lookup)

        if quickumls_path is not None and Path(quickumls_path).exists():
            self._quickumls, load_err = self._try_load_quickumls(quickumls_path)
            if quickumls_only and self._quickumls is None:
                raise RuntimeError(
                    "quickumls_only=True but QuickUMLS could not be loaded.\n"
                    f"{load_err or 'Unknown error.'}\n"
                    "Fix: run ./setup.sh in the bundle root, then use ./run_export.sh "
                    "(not system python3)."
                )

    @staticmethod
    def _try_load_quickumls(quickumls_path: Path):
        import sys

        try:
            from quickumls import QuickUMLS
        except ImportError as exc:
            return None, (
                f"ImportError: {exc}\n"
                f"  python: {sys.executable}\n"
                "  Hint: create the venv with ./setup.sh and run via ./run_export.sh"
            )
        try:
            matcher = QuickUMLS(
                str(quickumls_path),
                threshold=0.7,
                accepted_semtypes=sorted(CLINICAL_SEMTYPES),
            )
        except Exception as exc:
            return None, (
                f"{type(exc).__name__}: {exc}\n"
                f"  python: {sys.executable}\n"
                f"  quickumls_data: {quickumls_path}"
            )
        return matcher, None

    def _link_quickumls(self, text: str) -> list[tuple[str, float]]:
        if self._quickumls is None:
            return []
        out: list[tuple[str, float]] = []
        for match_group in self._quickumls.match(
            text, best_match=True, ignore_syntax=False
        ):
            for m in match_group:
                cui = m.get("cui")
                sim = float(m.get("similarity", 0.0))
                if cui and sim >= self.min_confidence:
                    out.append((cui, sim))
        return out

    def _link_dictionary(self, text: str) -> list[tuple[str, float]]:
        tokens = _tokenize(text)
        if not tokens:
            return []

        matches: dict[str, float] = {}
        n = len(tokens)
        for start in range(n):
            for length in range(min(self.max_ngram, n - start), 0, -1):
                phrase = " ".join(tokens[start : start + length])
                hits = self.surface_lookup.get(phrase)
                if not hits:
                    continue
                for cui, score in hits:
                    matches[cui] = max(matches.get(cui, 0.0), score)
                break  # longest match at this start position

        # Also try semicolon / clause segments (common in evidence strings).
        for clause in re.split(r"[;,\n]+", text):
            clause = clause.strip()
            if not clause:
                continue
            norm_clause = normalize_surface(clause)
            if norm_clause in self.surface_lookup:
                for cui, score in self.surface_lookup[norm_clause]:
                    matches[cui] = max(matches.get(cui, 0.0), score)

        return sorted(matches.items(), key=lambda x: -x[1])

    def link_cuis(self, text: str) -> list[tuple[str, float]]:
        """Return (cui, confidence) pairs for a query."""
        matches: dict[str, float] = {}
        if self.quickumls_only:
            for cui, score in self._link_quickumls(text):
                if self.skip_graph_cui_filter or cui in self.cui_to_idx:
                    matches[cui] = max(matches.get(cui, 0.0), score)
        else:
            for cui, score in self._link_dictionary(text):
                matches[cui] = max(matches.get(cui, 0.0), score)
            for cui, score in self._link_quickumls(text):
                if cui in self.cui_to_idx:
                    matches[cui] = max(matches.get(cui, 0.0), score)
        return sorted(matches.items(), key=lambda x: -x[1])

    def link_graph_nodes(self, text: str) -> list[tuple[int, float]]:
        """Return (graph_row_idx, confidence) pairs for a query."""
        out: list[tuple[int, float]] = []
        for cui, score in self.link_cuis(text):
            idx = self.cui_to_idx.get(cui)
            if idx is not None:
                out.append((int(idx), score))
        return out

    def coverage_stats(self, texts: list[str]) -> dict[str, float]:
        """Compute entity-linking coverage over a list of queries."""
        n = len(texts)
        if n == 0:
            return {
                "n": 0,
                "pct_with_match": 0.0,
                "pct_zero_match": 100.0,
                "avg_entities": 0.0,
            }
        counts = [len(self.link_graph_nodes(t)) for t in texts]
        with_match = sum(1 for c in counts if c > 0)
        return {
            "n": n,
            "pct_with_match": 100.0 * with_match / n,
            "pct_zero_match": 100.0 * (n - with_match) / n,
            "avg_entities": float(np.mean(counts)),
            "median_entities": float(np.median(counts)),
            "max_entities": int(max(counts)),
        }

	"""Link clinical evidence text to UMLS graph nodes.

	Primary backend: surface-form dictionary built from the filtered KG (``id_maps.parquet``)
	plus optional enrichment from ``mrxns_eng`` / ``mrsty`` in the UMLS DuckDB.

	Optional backend: QuickUMLS (when ``leveldb`` + ``quickumls`` are installed).
	"""

	from __future__ import annotations

	import json
	import re
	import unicodedata
	from pathlib import Path

	import numpy as np
	import polars as pl

	from config.paths import (
	CUI_TO_GRAPH_IDX_JSON,
	GRAPH_ID_MAPS_PARQUET,
	QUERY_LINKER_CACHE_JSON,
	UMLS_DUCKDB_PATH,
	)

	# Clinical semantic types (diseases, drugs, symptoms, procedures, anatomy, findings).
	CLINICAL_SEMTYPES: frozenset[str] = frozenset(
	{
	"T047", # Disease or Syndrome
	"T048", # Mental or Behavioral Dysfunction
	"T184", # Sign or Symptom
	"T121", # Pharmacologic Substance
	"T109", # Organic Chemical
	"T123", # Biologically Active Substance
	"T195", # Antibiotic
	"T200", # Clinical Drug
	"T023", # Body Part, Organ, or Organ Component
	"T031", # Body Substance
	"T033", # Finding
	"T034", # Laboratory or Test Result
	"T060", # Diagnostic Procedure
	"T061", # Therapeutic or Preventive Procedure
	"T037", # Injury or Poisoning
	"T046", # Pathologic Function
	"T191", # Neoplastic Process
	}
	)

	PREFERRED_SABS: frozenset[str] = frozenset(
	{
	"SNOMEDCT_US",
	"MSH",
	"RXNORM",
	"MEDCIN",
	"NCI",
	"LNC",
	"ICD10CM",
	"ICD10",
	"ICD9CM",
	"NDDF",
	"MMSL",
	"VANDF",
	}
	)

	STOPWORDS: frozenset[str] = frozenset(
	{
	"a",
	"an",
	"the",
	"and",
	"or",
	"of",
	"in",
	"on",
	"at",
	"to",
	"for",
	"with",
	"without",
	"by",
	"from",
	"is",
	"are",
	"was",
	"were",
	"be",
	"been",
	"being",
	"has",
	"have",
	"had",
	"not",
	"no",
	"due",
	"associated",
	"secondary",
	"primary",
	"unspecified",
	"other",
	"specified",
	"history",
	"personal",
	"patient",
	"given",
	"then",
	"there",
	"will",
	"some",
	"any",
	"all",
	"both",
	"each",
	"than",
	"that",
	"this",
	"these",
	"those",
	"into",
	"during",
	"after",
	"before",
	"while",
	"when",
	"where",
	"which",
	"who",
	"whom",
	"whose",
	"if",
	"as",
	"but",
	"so",
	"such",
	"via",
	"per",
	"also",
	"only",
	"just",
	"very",
	"more",
	"most",
	"less",
	"least",
	"new",
	"old",
	"acute",
	"chronic",
	"stage",
	"type",
	"status",
	"present",
	"absent",
	"possible",
	"probable",
	"suspected",
	}
	)

	_SURFACE_RE = re.compile(r"[^a-z0-9]+")


	def normalize_surface(text: str) -> str:
	"""Lowercase, strip accents, collapse punctuation to spaces."""
	text = unicodedata.normalize("NFKD", text)
	text = "".join(ch for ch in text if not unicodedata.combining(ch))
	text = text.lower().strip()
	text = _SURFACE_RE.sub(" ", text)
	return " ".join(text.split())


	def _tokenize(text: str) -> list[str]:
	return [t for t in normalize_surface(text).split() if t and t not in STOPWORDS]


	def build_cui_to_graph_idx(
	id_maps_parquet: Path = GRAPH_ID_MAPS_PARQUET,
	max_graph_idx: int \| None = None,
	) -> dict[str, int]:
	"""Map CUI -> best graph row index (highest degree node per CUI)."""
	scan = pl.scan_parquet(id_maps_parquet).select("idx", "cui", "degree")
	if max_graph_idx is not None:
	scan = scan.filter(pl.col("idx") < max_graph_idx)
	df = (
	scan.filter(pl.col("cui").is_not_null())
	.sort(["cui", "degree"], descending=[False, True])
	.group_by("cui")
	.agg(pl.col("idx").first())
	.collect()
	)
	return dict(zip(df["cui"].to_list(), df["idx"].to_list()))


	def build_clinical_cuis(
	cui_to_idx: dict[str, int],
	umls_duckdb: Path \| None = UMLS_DUCKDB_PATH,
	) -> set[str]:
	"""Filter CUIs to clinical semantic types when DuckDB is available."""
	if umls_duckdb is None or not Path(umls_duckdb).exists():
	return set(cui_to_idx)
	import duckdb

	con = duckdb.connect(str(umls_duckdb), read_only=True)
	con.register("graph_cuis", pl.DataFrame({"cui": list(cui_to_idx)}))
	sem = ",".join(f"'{s}'" for s in CLINICAL_SEMTYPES)
	rows = con.execute(
	f"""
	SELECT DISTINCT g.cui
	FROM graph_cuis g
	JOIN mrsty s ON g.cui = s.CUI
	WHERE s.STY IN ({sem})
	"""
	).fetchall()
	con.close()
	return {r[0] for r in rows} if rows else set(cui_to_idx)


	def build_surface_lookup(
	cui_to_idx: dict[str, int],
	clinical_cuis: set[str],
	id_maps_parquet: Path = GRAPH_ID_MAPS_PARQUET,
	umls_duckdb: Path \| None = UMLS_DUCKDB_PATH,
	max_graph_idx: int \| None = None,
	min_surface_len: int = 3,
	) -> dict[str, list[tuple[str, float]]]:
	"""Build normalized surface form -> [(cui, score)] from KG + optional MRXNS."""
	lookup: dict[str, dict[str, float]] = {}

	def add(surface: str, cui: str, score: float) -> None:
	if cui not in clinical_cuis:
	return
	norm = normalize_surface(surface)
	if len(norm) < min_surface_len:
	return
	bucket = lookup.setdefault(norm, {})
	bucket[cui] = max(bucket.get(cui, 0.0), score)

	scan = pl.scan_parquet(id_maps_parquet).select("idx", "cui", "sab", "str")
	if max_graph_idx is not None:
	scan = scan.filter(pl.col("idx") < max_graph_idx)
	df = scan.filter(
	pl.col("cui").is_not_null()
	& pl.col("str").is_not_null()
	& (pl.col("str").str.len_chars() >= min_surface_len)
	).collect()

	for cui, sab, surface in zip(
	df["cui"].to_list(), df["sab"].to_list(), df["str"].to_list()
	):
	score = 1.0 if sab in PREFERRED_SABS else 0.85
	add(surface, cui, score)

	if umls_duckdb is not None and Path(umls_duckdb).exists():
	import duckdb

	con = duckdb.connect(str(umls_duckdb), read_only=True)
	con.register("graph_cuis", pl.DataFrame({"cui": list(clinical_cuis)}))
	sem = ",".join(f"'{s}'" for s in CLINICAL_SEMTYPES)
	rows = con.execute(
	f"""
	SELECT DISTINCT m.CUI, m.NSTR
	FROM mrxns_eng m
	JOIN graph_cuis g ON m.CUI = g.cui
	JOIN mrsty s ON m.CUI = s.CUI
	WHERE s.STY IN ({sem})
	AND length(m.NSTR) >= {int(min_surface_len)}
	"""
	).fetchall()
	con.close()
	for cui, nstr in rows:
	surface = " ".join(nstr.split())
	add(surface, cui, 0.9)

	return {k: sorted(v.items(), key=lambda x: -x[1]) for k, v in lookup.items()}


	def save_cui_to_graph_idx(path: Path, cui_to_idx: dict[str, int]) -> None:
	path.parent.mkdir(parents=True, exist_ok=True)
	path.write_text(json.dumps(cui_to_idx), encoding="utf-8")


	def load_cui_to_graph_idx(path: Path = CUI_TO_GRAPH_IDX_JSON) -> dict[str, int]:
	payload = json.loads(path.read_text(encoding="utf-8"))
	return {str(k): int(v) for k, v in payload.items()}


	def ensure_cui_to_graph_idx(
	path: Path = CUI_TO_GRAPH_IDX_JSON,
	id_maps_parquet: Path = GRAPH_ID_MAPS_PARQUET,
	rebuild: bool = False,
	) -> dict[str, int]:
	"""Build or load CUI → graph row index (full graph, highest-degree node per CUI)."""
	if path.exists() and not rebuild:
	return load_cui_to_graph_idx(path)
	cui_to_idx = build_cui_to_graph_idx(id_maps_parquet, max_graph_idx=None)
	save_cui_to_graph_idx(path, cui_to_idx)
	return cui_to_idx


	def save_linker_cache(
	path: Path,
	cui_to_idx: dict[str, int],
	surface_lookup: dict[str, list[tuple[str, float]]],
	) -> None:
	path.parent.mkdir(parents=True, exist_ok=True)
	payload = {
	"cui_to_idx": cui_to_idx,
	"surface_lookup": {k: list(v) for k, v in surface_lookup.items()},
	}
	path.write_text(json.dumps(payload), encoding="utf-8")


	def load_linker_cache(
	path: Path,
	) -> tuple[dict[str, int], dict[str, list[tuple[str, float]]]]:
	payload = json.loads(path.read_text(encoding="utf-8"))
	surface_lookup = {
	k: [(cui, float(score)) for cui, score in v]
	for k, v in payload["surface_lookup"].items()
	}
	return payload["cui_to_idx"], surface_lookup


	class QueryEntityLinker:
	"""Map query text to graph node indices with confidence scores."""

	def __init__(
	self,
	*,
	id_maps_parquet: Path = GRAPH_ID_MAPS_PARQUET,
	umls_duckdb: Path \| None = UMLS_DUCKDB_PATH,
	cache_path: Path = QUERY_LINKER_CACHE_JSON,
	quickumls_path: Path \| None = None,
	max_graph_idx: int \| None = None,
	rebuild_cache: bool = False,
	min_confidence: float = 0.7,
	max_ngram: int = 6,
	quickumls_only: bool = False,
	skip_graph_cui_filter: bool = False,
	):
	self.min_confidence = min_confidence
	self.max_ngram = max_ngram
	self.quickumls_only = quickumls_only
	self.skip_graph_cui_filter = skip_graph_cui_filter
	self._quickumls = None

	if quickumls_only and skip_graph_cui_filter:
	self.cui_to_idx = {}
	self.surface_lookup = {}
	elif quickumls_only:
	self.cui_to_idx = ensure_cui_to_graph_idx(
	rebuild=rebuild_cache,
	id_maps_parquet=id_maps_parquet,
	)
	self.surface_lookup = {}
	elif cache_path.exists() and not rebuild_cache:
	self.cui_to_idx, self.surface_lookup = load_linker_cache(cache_path)
	else:
	# Build full-graph linker cache; OOB embedding rows are filtered at pool time.
	cui_to_idx = build_cui_to_graph_idx(id_maps_parquet, max_graph_idx=None)
	clinical_cuis = build_clinical_cuis(cui_to_idx, umls_duckdb)
	self.surface_lookup = build_surface_lookup(
	cui_to_idx,
	clinical_cuis,
	id_maps_parquet=id_maps_parquet,
	umls_duckdb=umls_duckdb,
	max_graph_idx=None,
	)
	self.cui_to_idx = {
	c: i for c, i in cui_to_idx.items() if c in clinical_cuis
	}
	save_linker_cache(cache_path, self.cui_to_idx, self.surface_lookup)

	if quickumls_path is not None and Path(quickumls_path).exists():
	self._quickumls, load_err = self._try_load_quickumls(quickumls_path)
	if quickumls_only and self._quickumls is None:
	raise RuntimeError(
	"quickumls_only=True but QuickUMLS could not be loaded.\n"
	f"{load_err or 'Unknown error.'}\n"
	"Fix: run ./setup.sh in the bundle root, then use ./run_export.sh "
	"(not system python3)."
	)

	@staticmethod
	def _try_load_quickumls(quickumls_path: Path):
	import sys

	try:
	from quickumls import QuickUMLS
	except ImportError as exc:
	return None, (
	f"ImportError: {exc}\n"
	f" python: {sys.executable}\n"
	" Hint: create the venv with ./setup.sh and run via ./run_export.sh"
	)
	try:
	matcher = QuickUMLS(
	str(quickumls_path),
	threshold=0.7,
	accepted_semtypes=sorted(CLINICAL_SEMTYPES),
	)
	except Exception as exc:
	return None, (
	f"{type(exc).__name__}: {exc}\n"
	f" python: {sys.executable}\n"
	f" quickumls_data: {quickumls_path}"
	)
	return matcher, None

	def _link_quickumls(self, text: str) -> list[tuple[str, float]]:
	if self._quickumls is None:
	return []
	out: list[tuple[str, float]] = []
	for match_group in self._quickumls.match(
	text, best_match=True, ignore_syntax=False
	):
	for m in match_group:
	cui = m.get("cui")
	sim = float(m.get("similarity", 0.0))
	if cui and sim >= self.min_confidence:
	out.append((cui, sim))
	return out

	def _link_dictionary(self, text: str) -> list[tuple[str, float]]:
	tokens = _tokenize(text)
	if not tokens:
	return []

	matches: dict[str, float] = {}
	n = len(tokens)
	for start in range(n):
	for length in range(min(self.max_ngram, n - start), 0, -1):
	phrase = " ".join(tokens[start : start + length])
	hits = self.surface_lookup.get(phrase)
	if not hits:
	continue
	for cui, score in hits:
	matches[cui] = max(matches.get(cui, 0.0), score)
	break # longest match at this start position

	# Also try semicolon / clause segments (common in evidence strings).
	for clause in re.split(r"[;,\n]+", text):
	clause = clause.strip()
	if not clause:
	continue
	norm_clause = normalize_surface(clause)
	if norm_clause in self.surface_lookup:
	for cui, score in self.surface_lookup[norm_clause]:
	matches[cui] = max(matches.get(cui, 0.0), score)

	return sorted(matches.items(), key=lambda x: -x[1])

	def link_cuis(self, text: str) -> list[tuple[str, float]]:
	"""Return (cui, confidence) pairs for a query."""
	matches: dict[str, float] = {}
	if self.quickumls_only:
	for cui, score in self._link_quickumls(text):
	if self.skip_graph_cui_filter or cui in self.cui_to_idx:
	matches[cui] = max(matches.get(cui, 0.0), score)
	else:
	for cui, score in self._link_dictionary(text):
	matches[cui] = max(matches.get(cui, 0.0), score)
	for cui, score in self._link_quickumls(text):
	if cui in self.cui_to_idx:
	matches[cui] = max(matches.get(cui, 0.0), score)
	return sorted(matches.items(), key=lambda x: -x[1])

	def link_graph_nodes(self, text: str) -> list[tuple[int, float]]:
	"""Return (graph_row_idx, confidence) pairs for a query."""
	out: list[tuple[int, float]] = []
	for cui, score in self.link_cuis(text):
	idx = self.cui_to_idx.get(cui)
	if idx is not None:
	out.append((int(idx), score))
	return out

	def coverage_stats(self, texts: list[str]) -> dict[str, float]:
	"""Compute entity-linking coverage over a list of queries."""
	n = len(texts)
	if n == 0:
	return {
	"n": 0,
	"pct_with_match": 0.0,
	"pct_zero_match": 100.0,
	"avg_entities": 0.0,
	}
	counts = [len(self.link_graph_nodes(t)) for t in texts]
	with_match = sum(1 for c in counts if c > 0)
	return {
	"n": n,
	"pct_with_match": 100.0 * with_match / n,
	"pct_zero_match": 100.0 * (n - with_match) / n,
	"avg_entities": float(np.mean(counts)),
	"median_entities": float(np.median(counts)),
	"max_entities": int(max(counts)),
	}