Spaces:
Runtime error
Runtime error
| """Export QuickUMLS-detected CUIs and clinical-AUI-validated AUIs for query texts.""" | |
| from __future__ import annotations | |
| import json | |
| from dataclasses import asdict, dataclass | |
| from pathlib import Path | |
| from typing import Iterable, Iterator | |
| from config.paths import ( | |
| CUI_TO_VALID_AUI_JSON, | |
| CUI_TO_VALID_AUI_PARQUET, | |
| GRAPH_ID_MAPS_PARQUET, | |
| MRCONSO_CLINICAL_PARQUET, | |
| MRCONSO_CLINICAL_PARQUET_DEFAULT, | |
| ) | |
| from src.data.query_entity_linker import QueryEntityLinker | |
| def _resolve_mrconso_parquet(path: Path | None) -> Path: | |
| if path is not None: | |
| candidate = path | |
| elif MRCONSO_CLINICAL_PARQUET.exists(): | |
| candidate = MRCONSO_CLINICAL_PARQUET | |
| else: | |
| candidate = MRCONSO_CLINICAL_PARQUET_DEFAULT | |
| if not candidate.exists(): | |
| raise FileNotFoundError( | |
| f"Clinical MRCONSO parquet not found: {candidate}. " | |
| "Copy mrconso_clinical.parquet to data/processed/graph/ or set " | |
| "ICD_MRCONSO_CLINICAL_PARQUET." | |
| ) | |
| return candidate | |
| class ValidAuiRecord: | |
| cui: str | |
| aui: str | |
| graph_idx: int | |
| str: str | |
| sab: str | |
| def _cui_to_valid_aui_sql( | |
| *, | |
| id_maps_parquet: Path, | |
| mrconso_parquet: Path, | |
| ) -> str: | |
| return f""" | |
| WITH graph_auis AS ( | |
| SELECT | |
| cui, | |
| REPLACE(node_id, 'AUI:', '') AS aui, | |
| idx, | |
| degree, | |
| str, | |
| sab | |
| FROM read_parquet('{id_maps_parquet.as_posix()}') | |
| WHERE kind = 'AUI' AND cui IS NOT NULL | |
| ), | |
| ranked AS ( | |
| SELECT | |
| g.cui, | |
| g.aui, | |
| g.idx, | |
| g.str, | |
| g.sab, | |
| ROW_NUMBER() OVER ( | |
| PARTITION BY g.cui | |
| ORDER BY g.degree DESC, g.idx | |
| ) AS rn | |
| FROM graph_auis g | |
| INNER JOIN read_parquet('{mrconso_parquet.as_posix()}') m | |
| ON g.aui = m.AUI | |
| ) | |
| SELECT cui, aui, idx AS graph_idx, str, sab | |
| FROM ranked | |
| WHERE rn = 1 | |
| """ | |
| def build_cui_to_valid_aui_parquet( | |
| output: Path = CUI_TO_VALID_AUI_PARQUET, | |
| *, | |
| id_maps_parquet: Path = GRAPH_ID_MAPS_PARQUET, | |
| mrconso_parquet: Path | None = None, | |
| ) -> int: | |
| """Build the CUI→AUI parquet cache via DuckDB + clinical MRCONSO parquet.""" | |
| mrconso = _resolve_mrconso_parquet(mrconso_parquet) | |
| if not id_maps_parquet.exists(): | |
| raise FileNotFoundError(f"Graph id maps not found: {id_maps_parquet}") | |
| import duckdb | |
| output.parent.mkdir(parents=True, exist_ok=True) | |
| con = duckdb.connect() | |
| con.execute( | |
| f""" | |
| COPY ( | |
| {_cui_to_valid_aui_sql( | |
| id_maps_parquet=id_maps_parquet, | |
| mrconso_parquet=mrconso, | |
| )} | |
| ) TO '{output.as_posix()}' (FORMAT PARQUET) | |
| """ | |
| ) | |
| con.close() | |
| import polars as pl | |
| return pl.scan_parquet(output).select(pl.len()).collect().item() | |
| def build_cui_to_valid_aui_map( | |
| *, | |
| id_maps_parquet: Path = GRAPH_ID_MAPS_PARQUET, | |
| mrconso_parquet: Path | None = None, | |
| ) -> dict[str, ValidAuiRecord]: | |
| """Map each CUI to its highest-degree AUI present in clinical MRCONSO.""" | |
| build_cui_to_valid_aui_parquet( | |
| id_maps_parquet=id_maps_parquet, | |
| mrconso_parquet=mrconso_parquet, | |
| ) | |
| lookup = CuiToValidAuiLookup(CUI_TO_VALID_AUI_PARQUET) | |
| try: | |
| import polars as pl | |
| cuis = ( | |
| pl.scan_parquet(CUI_TO_VALID_AUI_PARQUET) | |
| .select("cui") | |
| .collect() | |
| .get_column("cui") | |
| .to_list() | |
| ) | |
| return lookup.lookup_many(cuis) | |
| finally: | |
| lookup.close() | |
| def save_cui_to_valid_aui_map( | |
| path: Path, | |
| mapping: dict[str, ValidAuiRecord], | |
| ) -> None: | |
| path.parent.mkdir(parents=True, exist_ok=True) | |
| payload = {cui: asdict(rec) for cui, rec in mapping.items()} | |
| path.write_text(json.dumps(payload), encoding="utf-8") | |
| def load_cui_to_valid_aui_map( | |
| path: Path = CUI_TO_VALID_AUI_JSON, | |
| ) -> dict[str, ValidAuiRecord]: | |
| payload = json.loads(path.read_text(encoding="utf-8")) | |
| return { | |
| cui: ValidAuiRecord( | |
| cui=str(rec["cui"]), | |
| aui=str(rec["aui"]), | |
| graph_idx=int(rec["graph_idx"]), | |
| str=str(rec.get("str", "")), | |
| sab=str(rec.get("sab", "")), | |
| ) | |
| for cui, rec in payload.items() | |
| } | |
| class CuiToValidAuiLookup: | |
| """Polars-backed lookup over ``cui_to_valid_aui.parquet`` (memory-safe).""" | |
| def __init__(self, parquet_path: Path = CUI_TO_VALID_AUI_PARQUET) -> None: | |
| if not parquet_path.exists(): | |
| raise FileNotFoundError(f"CUI→AUI parquet not found: {parquet_path}") | |
| import polars as pl | |
| self._parquet_path = parquet_path | |
| self._scan = pl.scan_parquet(parquet_path) | |
| def count(self) -> int: | |
| import polars as pl | |
| return self._scan.select(pl.len()).collect().item() | |
| def lookup_many(self, cuis: Iterable[str]) -> dict[str, ValidAuiRecord]: | |
| unique = list(dict.fromkeys(str(c) for c in cuis if c)) | |
| if not unique: | |
| return {} | |
| import polars as pl | |
| df = self._scan.filter(pl.col("cui").is_in(unique)).collect() | |
| out: dict[str, ValidAuiRecord] = {} | |
| for row in df.iter_rows(named=True): | |
| cui = str(row["cui"]) | |
| out[cui] = ValidAuiRecord( | |
| cui=cui, | |
| aui=str(row["aui"]), | |
| graph_idx=int(row["graph_idx"]), | |
| str=str(row["str"]) if row["str"] is not None else "", | |
| sab=str(row["sab"]) if row["sab"] is not None else "", | |
| ) | |
| return out | |
| def close(self) -> None: | |
| return None | |
| def ensure_cui_to_valid_aui_lookup( | |
| *, | |
| parquet_path: Path = CUI_TO_VALID_AUI_PARQUET, | |
| rebuild: bool = False, | |
| id_maps_parquet: Path = GRAPH_ID_MAPS_PARQUET, | |
| mrconso_parquet: Path | None = None, | |
| ) -> CuiToValidAuiLookup: | |
| if not parquet_path.exists() or rebuild: | |
| build_cui_to_valid_aui_parquet( | |
| parquet_path, | |
| id_maps_parquet=id_maps_parquet, | |
| mrconso_parquet=mrconso_parquet, | |
| ) | |
| return CuiToValidAuiLookup(parquet_path) | |
| def ensure_cui_to_valid_aui_map( | |
| *, | |
| path: Path = CUI_TO_VALID_AUI_JSON, | |
| rebuild: bool = False, | |
| id_maps_parquet: Path = GRAPH_ID_MAPS_PARQUET, | |
| mrconso_parquet: Path | None = None, | |
| ) -> dict[str, ValidAuiRecord]: | |
| lookup = ensure_cui_to_valid_aui_lookup( | |
| rebuild=rebuild, | |
| id_maps_parquet=id_maps_parquet, | |
| mrconso_parquet=mrconso_parquet, | |
| ) | |
| try: | |
| import polars as pl | |
| cuis = ( | |
| pl.scan_parquet(CUI_TO_VALID_AUI_PARQUET) | |
| .select("cui") | |
| .collect() | |
| .get_column("cui") | |
| .to_list() | |
| ) | |
| return lookup.lookup_many(cuis) | |
| finally: | |
| lookup.close() | |
| def detect_entities_for_text( | |
| text: str, | |
| *, | |
| linker: QueryEntityLinker, | |
| lookup: CuiToValidAuiLookup, | |
| ) -> list[dict[str, object]]: | |
| """Return kept CUI/AUI entities for one query text.""" | |
| linked = linker.link_cuis(text) | |
| records = lookup.lookup_many(cui for cui, _ in linked) | |
| entities: list[dict[str, object]] = [] | |
| for cui, confidence in linked: | |
| record = records.get(cui) | |
| if record is None: | |
| continue | |
| entities.append( | |
| { | |
| "cui": record.cui, | |
| "aui": record.aui, | |
| "confidence": round(float(confidence), 4), | |
| "graph_idx": record.graph_idx, | |
| "str": record.str, | |
| "sab": record.sab, | |
| } | |
| ) | |
| return entities | |
| def export_entity_detections( | |
| rows: list[dict[str, object]], | |
| *, | |
| linker: QueryEntityLinker, | |
| lookup: CuiToValidAuiLookup, | |
| text_key: str = "query_text", | |
| include_code: bool = True, | |
| ) -> list[dict[str, object]]: | |
| """Build export records aligned to input JSONL rows.""" | |
| return list( | |
| iter_entity_detection_entries( | |
| rows, | |
| linker=linker, | |
| lookup=lookup, | |
| text_key=text_key, | |
| include_code=include_code, | |
| ) | |
| ) | |
| def iter_entity_detection_entries( | |
| rows: list[dict[str, object]], | |
| *, | |
| linker: QueryEntityLinker, | |
| lookup: CuiToValidAuiLookup, | |
| text_key: str = "query_text", | |
| include_code: bool = True, | |
| ) -> Iterator[dict[str, object]]: | |
| for row in rows: | |
| text = str(row.get(text_key, "")).strip() | |
| entry: dict[str, object] = {"query_text": text} | |
| if include_code and "code" in row: | |
| entry["code"] = row["code"] | |
| entry["entities"] = detect_entities_for_text( | |
| text, | |
| linker=linker, | |
| lookup=lookup, | |
| ) | |
| yield entry | |
| def write_entity_detections_json( | |
| entries: Iterable[dict[str, object]], | |
| output: Path, | |
| *, | |
| pretty: bool = False, | |
| ) -> tuple[int, int, int]: | |
| """Stream entries to JSON without holding the full payload in memory.""" | |
| output.parent.mkdir(parents=True, exist_ok=True) | |
| n_entries = 0 | |
| with_entities = 0 | |
| total_entities = 0 | |
| with open(output, "w", encoding="utf-8") as f: | |
| f.write('{\n "entries": [\n') | |
| first = True | |
| for entry in entries: | |
| if not first: | |
| f.write(",\n") | |
| else: | |
| first = False | |
| f.write(" ") | |
| if pretty: | |
| rendered = json.dumps(entry, ensure_ascii=False, indent=2) | |
| f.write(rendered.replace("\n", "\n ")) | |
| else: | |
| f.write(json.dumps(entry, ensure_ascii=False)) | |
| n_entries += 1 | |
| n_ent = len(entry.get("entities", [])) | |
| if n_ent: | |
| with_entities += 1 | |
| total_entities += n_ent | |
| f.write("\n ]\n}\n") | |
| return n_entries, with_entities, total_entities | |