scripts/build_pinecone_index.py

"""Offline (v2): embed the MSigDB corpus by concept name and upsert to Pinecone.

This is the one-time build for the optional semantic-search / annotation-matching module. It needs
`PINECONE_API_KEY` and downloads the MSigDB metadata JSON once. Staged-subset default
(Hallmark + Reactome + GO:BP, ~9.4K sets); widen `--collections` to scale toward the full corpus.

    PINECONE_API_KEY=... python scripts/build_pinecone_index.py                 # staged subset
    PINECONE_API_KEY=... python scripts/build_pinecone_index.py --limit 500      # quick smoke
    PINECONE_API_KEY=... python scripts/build_pinecone_index.py --collections H C2:CP:REACTOME C5:GO:BP C5:GO:MF
"""
from __future__ import annotations

import argparse
import json
import os
import sys
import time
import urllib.request
from pathlib import Path

import numpy as np

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))

import config  # noqa: E402
from mcp_server.corpus import DEFAULT_COLLECTIONS, concept_text  # noqa: E402

MSIGDB_URL = "https://data.broadinstitute.org/gsea-msigdb/msigdb/release/2024.1.Hs/msigdb.v2024.1.Hs.json"
CACHE_JSON = Path(os.environ.get("SEMBLANCE_MSIGDB_JSON", "/tmp/msigdb.json"))


def load_corpus(collections: set[str], json_path: Path, limit: int | None) -> list[dict]:
    if not json_path.exists():
        print(f"downloading MSigDB JSON → {json_path} …")
        urllib.request.urlretrieve(MSIGDB_URL, json_path)
    data = json.loads(json_path.read_text())
    records = []
    for name, meta in data.items():
        coll = meta.get("collection")
        if coll not in collections:
            continue
        records.append({"name": name, "collection": coll, "concept": concept_text(name)})
        if limit and len(records) >= limit:
            break
    return records


def ensure_index(pc, name: str, dim: int):
    from pinecone import ServerlessSpec

    existing = {i.name for i in pc.list_indexes()}
    if name not in existing:
        print(f"creating serverless index '{name}' ({config.PINECONE_CLOUD}/{config.PINECONE_REGION}) …")
        pc.create_index(name=name, dimension=dim, metric="cosine",
                        spec=ServerlessSpec(cloud=config.PINECONE_CLOUD, region=config.PINECONE_REGION))
        for _ in range(60):
            if pc.describe_index(name).status.get("ready"):
                break
            time.sleep(1)
    return pc.Index(name)


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--collections", nargs="+", default=list(DEFAULT_COLLECTIONS))
    ap.add_argument("--limit", type=int, default=None, help="cap total sets (for a quick smoke)")
    ap.add_argument("--embed-batch", type=int, default=256)
    ap.add_argument("--upsert-batch", type=int, default=200)
    ap.add_argument("--json-path", type=Path, default=CACHE_JSON)
    ap.add_argument("--index", default=config.PINECONE_INDEX)
    ap.add_argument("--recreate", action="store_true", help="delete and rebuild the index first")
    args = ap.parse_args()

    api_key = os.environ.get("PINECONE_API_KEY")
    if not api_key:
        sys.exit("PINECONE_API_KEY is not set.")

    records = load_corpus(set(args.collections), args.json_path, args.limit)
    print(f"{len(records)} sets from {args.collections}")
    if not records:
        sys.exit("no records matched the requested collections")

    from pinecone import Pinecone
    from sentence_transformers import SentenceTransformer

    pc = Pinecone(api_key=api_key)
    if args.recreate and args.index in {i.name for i in pc.list_indexes()}:
        print(f"deleting existing index '{args.index}' …")
        pc.delete_index(args.index)
    index = ensure_index(pc, args.index, config.EMBED_DIM)

    print(f"loading BioLORD ({config.EMBEDDER_DEFAULT}) …")
    model = SentenceTransformer(config.EMBEDDER_DEFAULT, device=config.EMBED_DEVICE)

    concepts = [r["concept"] for r in records]
    upserted = 0
    for start in range(0, len(records), args.embed_batch):
        chunk = records[start : start + args.embed_batch]
        vecs = model.encode([r["concept"] for r in chunk], normalize_embeddings=True,
                            batch_size=64, show_progress_bar=False)
        items = [
            {"id": r["name"], "values": np.asarray(v, dtype=float).tolist(),
             "metadata": {"name": r["name"], "collection": r["collection"], "concept": r["concept"]}}
            for r, v in zip(chunk, vecs)
        ]
        for u in range(0, len(items), args.upsert_batch):
            index.upsert(vectors=items[u : u + args.upsert_batch])
        upserted += len(items)
        print(f"  upserted {upserted}/{len(records)}")

    stats = index.describe_index_stats()
    print(f"\ndone. index '{args.index}' vector count: {stats.get('total_vector_count')}")

    # quick self-check query
    q = model.encode(["interferon signaling"], normalize_embeddings=True)[0]
    res = index.query(vector=np.asarray(q, dtype=float).tolist(), top_k=5, include_metadata=True)
    print("sample query 'interferon signaling' →")
    for m in res.get("matches", []):
        md = m.get("metadata", {})
        print(f"  {m['score']:.3f}  {md.get('name')}  [{md.get('collection')}]")


if __name__ == "__main__":
    main()