build_semantic_indices.py

"""
build_semantic_indices.py
=========================
Pre-compute MedEmbed embeddings for semantic fallback matching.

Generates three embedding sets:
  1. Symptom canonical names   → hpo_data_final/symptom_embeddings.npy + .json
  2. MONDO disease synonyms    → hpo_data_final/disease_embeddings.npy + .json
  3. KG disease names          → hpo_data_final/kg_disease_embeddings.npy + .json

Each .json index is a list of [id, name] in row-order of the .npy array.
"""

import json
import numpy as np
from pathlib import Path
from sentence_transformers import SentenceTransformer

MONDO_DIR = Path(__file__).parent / "mondo_data_new"
OUT_DIR   = Path(__file__).parent / "hpo_data_final"
OUT_DIR.mkdir(parents=True, exist_ok=True)

print("Loading MedEmbed model...")
model = SentenceTransformer("abhinand/MedEmbed-large-v0.1", device="cpu")

# ---------------------------------------------------------------------------
# 1. Symptom canonical names (skip if already built)
# ---------------------------------------------------------------------------
sym_npy  = OUT_DIR / "symptom_embeddings.npy"
sym_json = OUT_DIR / "symptom_embedding_index.json"
if sym_npy.exists() and sym_json.exists():
    print("\n[1/3] Symptom embeddings already exist, skipping.")
else:
    print("\n[1/3] Symptom canonical names...")
    sc = json.load(open(MONDO_DIR / "symptom_catalog.json"))
    id_to_symptom = {hid: entry["label"] for hid, entry in sc.items()}
    sym_labels = [(hid, name) for hid, name in id_to_symptom.items()]
    sym_names  = [name for _, name in sym_labels]
    print(f"  Encoding {len(sym_names):,} names...")
    sym_emb = model.encode(sym_names, normalize_embeddings=True, show_progress_bar=True)
    np.save(sym_npy, sym_emb)
    with open(sym_json, "w") as f:
        json.dump(sym_labels, f)
    print(f"  Saved symptom_embeddings.npy  ({sym_emb.shape})")
    print(f"  Saved symptom_embedding_index.json  ({len(sym_labels)} entries)")

# ---------------------------------------------------------------------------
# 2. MONDO disease canonical labels (skip if already built)
# ---------------------------------------------------------------------------
dis_npy  = OUT_DIR / "disease_embeddings.npy"
dis_json = OUT_DIR / "disease_embedding_index.json"
if dis_npy.exists() and dis_json.exists():
    print("\n[2/3] MONDO disease embeddings already exist, skipping.")
else:
    print("\n[2/3] MONDO disease canonical labels...")
    dc = json.load(open(MONDO_DIR / "disease_catalog.json"))
    dis_labels = [(did, entry["label"]) for did, entry in dc.items()]
    dis_names  = [name for _, name in dis_labels]
    print(f"  Encoding {len(dis_names):,} names...")
    dis_emb = model.encode(dis_names, normalize_embeddings=True, show_progress_bar=True)
    np.save(dis_npy, dis_emb)
    with open(dis_json, "w") as f:
        json.dump(dis_labels, f)
    print(f"  Saved disease_embeddings.npy  ({dis_emb.shape})")
    print(f"  Saved disease_embedding_index.json  ({len(dis_labels)} entries)")

# ---------------------------------------------------------------------------
# 3. KG disease names — only diseases with scorable (HPO-bridgeable) phenotypes
# ---------------------------------------------------------------------------
kg_npy  = OUT_DIR / "kg_disease_embeddings.npy"
kg_json = OUT_DIR / "kg_disease_embedding_index.json"
if kg_npy.exists() and kg_json.exists():
    print("\n[3/3] KG disease embeddings already exist, skipping.")
else:
    print("\n[3/3] KG disease names (scorable only)...")

    # Load kg.feather to build phenotype x_id → symptom index
    import pandas as pd
    kg_df = pd.read_feather(str(Path(__file__).parent / "kg.feather"))
    pheno = kg_df[kg_df["relation"] == "disease_phenotype_positive"]

    # Build parent/child map for dedup (same as scanner)
    pp = kg_df[kg_df["relation"] == "phenotype_phenotype"]
    parent_of: dict[str, str] = {}
    for _, row in pp.iterrows():
        parent_of[str(row["y_name"]).strip().lower()] = str(row["x_name"]).strip().lower()

    # Load canonical symptom name → HPO ID mapping for bridging
    sc = json.load(open(MONDO_DIR / "symptom_catalog.json"))
    canon_name_to_hpo = {entry["label"].lower(): hid for hid, entry in sc.items()}

    # Group by x_id: dedup symptoms, then bridge to HPO IDs
    kg_labels = []
    for xid, group in pheno.groupby("x_id"):
        disease_name = str(group["x_name"].iloc[0]).strip()
        if not disease_name:
            continue
        raw = set(str(s).strip() for s in group["y_name"].tolist())
        # Dedup (same logic as scanner._dedup_symptoms)
        lower_map = {s.lower(): s for s in raw}
        deduped = set()
        for s_lower, s_orig in lower_map.items():
            if s_lower in parent_of and parent_of[s_lower] in lower_map:
                continue
            deduped.add(s_orig)
        # Bridge to HPO IDs
        bridged = {canon_name_to_hpo.get(s.lower()) for s in deduped if canon_name_to_hpo.get(s.lower())}
        if bridged:
            kg_labels.append((xid, disease_name))

    print(f"  {len(kg_labels):,} scorable KG x_ids")

    if kg_labels:
        kg_names = [name for _, name in kg_labels]
        print(f"  Encoding {len(kg_names):,} names...")
        kg_emb = model.encode(kg_names, normalize_embeddings=True, show_progress_bar=True)
        np.save(kg_npy, kg_emb)
        with open(kg_json, "w") as f:
            json.dump(kg_labels, f)
        print(f"  Saved kg_disease_embeddings.npy  ({kg_emb.shape})")
        print(f"  Saved kg_disease_embedding_index.json  ({len(kg_labels)} entries)")
    else:
        print("  WARNING: no KG diseases with scorable phenotypes found!")
        np.save(kg_npy, np.zeros((0, 1024), dtype="float32"))
        with open(kg_json, "w") as f:
            json.dump([], f)

print("\nDone — all semantic indices built.")