scripts/export_query_entity_detections.py

#!/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()