Spaces:
Runtime error
Runtime error
| #!/usr/bin/env python3 | |
| """Export QuickUMLS-detected CUIs and DuckDB-validated AUIs for query texts. | |
| Run with the QuickUMLS venv (required for entity detection):: | |
| /path/to/.venv-quickumls/bin/python3 scripts/export_query_entity_detections.py | |
| Single query:: | |
| .../python3 scripts/export_query_entity_detections.py \\ | |
| --query "Type II Diabetes Mellitus Uncontrolled" | |
| Full test split (writes ``results/query_entity_detections_test.json``):: | |
| .../python3 scripts/export_query_entity_detections.py --split test | |
| Build only the CUI→AUI cache (project venv is fine):: | |
| python3 scripts/export_query_entity_detections.py --build-aui-map-only | |
| """ | |
| from __future__ import annotations | |
| import argparse | |
| import sys | |
| from pathlib import Path | |
| WORKSPACE = Path(__file__).resolve().parent.parent | |
| if str(WORKSPACE) not in sys.path: | |
| sys.path.insert(0, str(WORKSPACE)) | |
| from config.paths import ( | |
| CUI_TO_VALID_AUI_PARQUET, | |
| QUICKUMLS_DATA_DIR, | |
| RESULTS_DIR, | |
| TEST_JSONL, | |
| TRAIN_JSONL, | |
| VAL_JSONL, | |
| ensure_dirs, | |
| ) | |
| from src.data.query_entity_export import ( | |
| build_cui_to_valid_aui_parquet, | |
| ensure_cui_to_valid_aui_lookup, | |
| iter_entity_detection_entries, | |
| write_entity_detections_json, | |
| ) | |
| from src.data.query_entity_linker import QueryEntityLinker | |
| SPLIT_PATHS = { | |
| "train": TRAIN_JSONL, | |
| "val": VAL_JSONL, | |
| "test": TEST_JSONL, | |
| } | |
| def _read_jsonl(path: Path) -> list[dict]: | |
| import polars as pl | |
| return pl.read_ndjson(path).to_dicts() | |
| def parse_args() -> argparse.Namespace: | |
| p = argparse.ArgumentParser( | |
| description="Export QuickUMLS CUIs and DuckDB-validated AUIs for queries." | |
| ) | |
| p.add_argument( | |
| "--query", | |
| type=str, | |
| default=None, | |
| help="Single query text to process.", | |
| ) | |
| p.add_argument( | |
| "--split", | |
| choices=sorted(SPLIT_PATHS), | |
| default=None, | |
| help="Process a train/val/test JSONL split.", | |
| ) | |
| p.add_argument( | |
| "--jsonl", | |
| type=Path, | |
| default=None, | |
| help="Custom JSONL input (uses query_text field).", | |
| ) | |
| p.add_argument( | |
| "--output", | |
| type=Path, | |
| default=None, | |
| help="Output JSON path (default: results/query_entity_detections_<name>.json).", | |
| ) | |
| p.add_argument( | |
| "--limit", | |
| type=int, | |
| default=None, | |
| help="Process only the first N rows.", | |
| ) | |
| p.add_argument( | |
| "--rebuild-aui-map", | |
| action="store_true", | |
| help="Rebuild cui_to_valid_aui.parquet from id_maps + DuckDB.", | |
| ) | |
| p.add_argument( | |
| "--build-aui-map-only", | |
| action="store_true", | |
| help="Only build/load the CUI→AUI parquet cache and exit.", | |
| ) | |
| p.add_argument( | |
| "--pretty", | |
| action="store_true", | |
| help="Pretty-print output JSON (uses more disk; default is compact).", | |
| ) | |
| p.add_argument( | |
| "--quickumls-dir", | |
| type=Path, | |
| default=QUICKUMLS_DATA_DIR, | |
| help="QuickUMLS data directory.", | |
| ) | |
| return p.parse_args() | |
| def main() -> None: | |
| args = parse_args() | |
| ensure_dirs() | |
| if args.build_aui_map_only: | |
| n_cuis = build_cui_to_valid_aui_parquet() | |
| print( | |
| f"Built CUI→valid AUI parquet: {n_cuis} CUIs -> {CUI_TO_VALID_AUI_PARQUET}" | |
| ) | |
| return | |
| if args.query: | |
| rows = [{"query_text": args.query}] | |
| default_name = "single" | |
| elif args.jsonl is not None: | |
| rows = _read_jsonl(args.jsonl) | |
| default_name = args.jsonl.stem | |
| elif args.split is not None: | |
| rows = _read_jsonl(SPLIT_PATHS[args.split]) | |
| default_name = args.split | |
| else: | |
| rows = _read_jsonl(TEST_JSONL) | |
| default_name = "test" | |
| if args.limit is not None: | |
| rows = rows[: args.limit] | |
| if not args.quickumls_dir.exists(): | |
| raise FileNotFoundError( | |
| f"QuickUMLS data not found at {args.quickumls_dir}. " | |
| "Set --quickumls-dir or ICD_QUICKUMLS_DATA_DIR." | |
| ) | |
| print("Loading CUI→valid AUI lookup (Polars over parquet)...") | |
| lookup = ensure_cui_to_valid_aui_lookup(rebuild=args.rebuild_aui_map) | |
| print(f" {lookup.count()} CUIs with a valid clinical AUI in mrconso_clinical.parquet") | |
| print(f"Loading QuickUMLS linker from {args.quickumls_dir} ...") | |
| linker = QueryEntityLinker( | |
| quickumls_path=args.quickumls_dir, | |
| quickumls_only=True, | |
| skip_graph_cui_filter=True, | |
| ) | |
| output = args.output or ( | |
| RESULTS_DIR / f"query_entity_detections_{default_name}.json" | |
| ) | |
| print(f"Processing {len(rows)} queries -> {output}") | |
| try: | |
| entries = iter_entity_detection_entries( | |
| rows, | |
| linker=linker, | |
| lookup=lookup, | |
| ) | |
| n_entries, with_entities, total_entities = write_entity_detections_json( | |
| entries, | |
| output, | |
| pretty=args.pretty, | |
| ) | |
| finally: | |
| lookup.close() | |
| print(f"Wrote {output}") | |
| print( | |
| f" queries with entities: {with_entities}/{n_entries} " | |
| f"(total entities kept: {total_entities})" | |
| ) | |
| if __name__ == "__main__": | |
| main() | |