Spaces:

Ruurd
/

semantic-harmonization-api

Sleeping

semantic-harmonization-api / src /retrieval /embeddings.py

Deploy OpenAI classification and protected terminology streaming

c2b7494 verified 5 days ago

20.4 kB

"""Embedding-based retrieval using local or API-backed encoders."""

from __future__ import annotations

from collections.abc import Callable
from dataclasses import dataclass
import os
from pathlib import Path
import re
from typing import Protocol, Sequence
from urllib.error import URLError
from urllib.request import urlopen

import numpy as np
import pandas as pd

DEFAULT_LOCAL_EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
DEFAULT_OPENAI_EMBEDDING_MODEL = "text-embedding-3-large"


def _create_local_embedding_model(model_name: str):
    """Import the optional local stack only when that backend is selected."""
    from sentence_transformers import SentenceTransformer

    return SentenceTransformer(model_name)


class EmbeddingEncoder(Protocol):
    """Common interface for local and API-backed embedding models."""

    def encode(self, sentences, **kwargs): ...


class OpenAIEmbeddingModel:
    """SentenceTransformer-compatible adapter for OpenAI embeddings."""

    def __init__(
        self,
        model_name: str = DEFAULT_OPENAI_EMBEDDING_MODEL,
        *,
        dimensions: int | None = 3072,
        batch_size: int = 512,
    ) -> None:
        if dimensions is not None and dimensions < 1:
            raise ValueError("Embedding dimensions must be at least 1")
        if batch_size < 1:
            raise ValueError("Embedding batch size must be at least 1")
        self.model_name = model_name
        self.dimensions = dimensions
        self.batch_size = batch_size
        self._client = None

    def _get_client(self):
        if self._client is None:
            import openai

            self._client = openai.OpenAI()
        return self._client

    def encode(
        self,
        sentences,
        *,
        normalize_embeddings: bool = True,
        show_progress_bar: bool = False,
        **_kwargs,
    ) -> np.ndarray:
        del show_progress_bar
        single = isinstance(sentences, str)
        inputs = [sentences] if single else [str(item) for item in sentences]
        if not inputs:
            width = self.dimensions or 0
            return np.empty((0, width), dtype=np.float32)

        # Convert each response to float32 immediately. Keeping hundreds of
        # thousands of embeddings as Python float lists can require many times
        # more memory than the final NumPy matrix.
        chunks: list[np.ndarray] = []
        client = self._get_client()
        for start in range(0, len(inputs), self.batch_size):
            request = {
                "model": self.model_name,
                "input": inputs[start : start + self.batch_size],
                "encoding_format": "float",
            }
            if self.dimensions is not None:
                request["dimensions"] = self.dimensions
            response = client.embeddings.create(**request)
            chunks.append(
                np.asarray(
                    [
                        item.embedding
                        for item in sorted(response.data, key=lambda item: item.index)
                    ],
                    dtype=np.float32,
                )
            )

        result = np.vstack(chunks)
        if normalize_embeddings:
            norms = np.linalg.norm(result, axis=1, keepdims=True)
            result = result / np.maximum(norms, 1e-12)
        return result[0] if single else result


def create_embedding_model(
    provider: str,
    model_name: str,
    *,
    dimensions: int | None = None,
    batch_size: int = 512,
) -> EmbeddingEncoder:
    """Construct a local SentenceTransformer or OpenAI embedding adapter."""
    provider = provider.strip().lower()
    if provider == "local":
        return _create_local_embedding_model(model_name)
    if provider == "openai":
        return OpenAIEmbeddingModel(
            model_name,
            dimensions=dimensions,
            batch_size=batch_size,
        )
    raise ValueError(f"Unsupported embedding provider: {provider!r}")


def embedding_cache_key(
    provider: str,
    model_name: str,
    dimensions: int | None = None,
) -> str:
    """Return a filesystem-safe cache identity without mixing providers/models."""
    model_key = re.sub(r"[^A-Za-z0-9_.-]+", "_", model_name).strip("_")
    if provider.strip().lower() == "local":
        return model_key
    dimension_key = f"_{dimensions}" if dimensions is not None else ""
    return f"{provider.strip().lower()}_{model_key}{dimension_key}"


@dataclass(slots=True)
class EmbeddingCandidate:
    rank: int
    code: str
    description: str
    vocabulary: str
    score: float

    def __lt__(self, other: object) -> bool:
        if not isinstance(other, EmbeddingCandidate):
            return NotImplemented
        return self.score > other.score


@dataclass
class EmbeddingIndex:
    """Wraps embedding model + corpus embeddings for fast cosine-sim search."""

    model: EmbeddingEncoder | None
    descriptions: list[str]
    codes: list[str]
    vocabularies: list[str]
    embeddings: np.ndarray  # shape (n, dim)

    @classmethod
    def from_codes(
        cls,
        codes: pd.DataFrame,
        model_name: str = DEFAULT_LOCAL_EMBEDDING_MODEL,
        model: EmbeddingEncoder | None = None,
        show_progress_bar: bool = False,
    ) -> "EmbeddingIndex":
        df = codes.copy()
        for col in ("code", "vocabulary"):
            if col not in df.columns:
                df[col] = ""
        embedding_model = model or _create_local_embedding_model(model_name)
        descs = df["description"].fillna("").astype(str).tolist()
        embs = embedding_model.encode(
            descs,
            show_progress_bar=show_progress_bar,
            normalize_embeddings=True,
        )
        return cls(
            model=embedding_model,
            descriptions=descs,
            codes=df["code"].astype(str).tolist(),
            vocabularies=df["vocabulary"].astype(str).tolist(),
            embeddings=np.asarray(embs, dtype=np.float32),
        )

    def save(self, path: str | Path) -> None:
        """Persist the corpus embeddings + metadata to disk (npz) so they can
        be reloaded without re-running the embedding model."""
        path = Path(path)
        path.parent.mkdir(parents=True, exist_ok=True)
        np.savez_compressed(
            path,
            embeddings=np.asarray(self.embeddings, dtype=np.float32),
            codes=np.array(self.codes, dtype=object),
            descriptions=np.array(self.descriptions, dtype=object),
            vocabularies=np.array(self.vocabularies, dtype=object),
        )

    @classmethod
    def load(
        cls,
        path: str | Path,
        model_name: str = DEFAULT_LOCAL_EMBEDDING_MODEL,
        model: EmbeddingEncoder | None = None,
        load_model: bool = True,
    ) -> "EmbeddingIndex":
        """Load a previously-saved embedding index from disk."""
        with np.load(path, allow_pickle=True) as data:
            return cls(
                model=model or (_create_local_embedding_model(model_name) if load_model else None),
                descriptions=list(data["descriptions"]),
                codes=list(data["codes"]),
                vocabularies=list(data["vocabularies"]),
                embeddings=np.asarray(data["embeddings"], dtype=np.float32),
            )

    @classmethod
    def from_cache_or_build(
        cls,
        codes: pd.DataFrame,
        cache_path: str | Path,
        model_name: str = DEFAULT_LOCAL_EMBEDDING_MODEL,
        model: EmbeddingEncoder | None = None,
        load_model: bool = True,
        show_progress_bar: bool = False,
        status_callback: Callable[[str], None] | None = None,
    ) -> "EmbeddingIndex":
        """Load a cached embedding index if present, otherwise build it from
        *codes* and cache the result for subsequent runs.
        This is the key mechanism that avoids re-embedding large reference
        code systems (e.g. the ~74k-row full ICD-10-CM corpus) on every run.
        """
        cache_path = Path(cache_path)
        if cache_path.exists():
            cached = cls.load(
                cache_path,
                model_name=model_name,
                model=model,
                load_model=load_model,
            )
            expected_codes = codes.get("code", pd.Series(dtype=str)).astype(str).tolist()
            expected_vocabularies = (
                codes.get("vocabulary", pd.Series(dtype=str)).astype(str).tolist()
            )
            expected_descriptions = codes["description"].fillna("").astype(str).tolist()
            if (
                cached.codes == expected_codes
                and cached.vocabularies == expected_vocabularies
                and cached.descriptions == expected_descriptions
            ):
                return cached
            subset_positions = _ordered_subset_positions(
                cached.codes,
                cached.vocabularies,
                expected_codes,
                expected_vocabularies,
            )
            if subset_positions is not None and all(
                cached.descriptions[position] == description
                for position, description in zip(subset_positions, expected_descriptions)
            ):
                # Corpus preprocessing can intentionally remove non-codable
                # dictionary rows (for example ICD range headers). Reuse the
                # matching cached vectors instead of re-embedding every code.
                subset = cls(
                    model=cached.model,
                    descriptions=expected_descriptions,
                    codes=expected_codes,
                    vocabularies=expected_vocabularies,
                    embeddings=np.asarray(cached.embeddings[subset_positions], dtype=np.float32),
                )
                subset.save(cache_path)
                return subset
        if status_callback is not None:
            status_callback(
                f"Creating embeddings for {len(codes):,} descriptions; no matching saved index. "
                "The first run can take several minutes."
            )
        index = cls.from_codes(
            codes,
            model_name=model_name,
            model=model,
            show_progress_bar=show_progress_bar,
        )
        index.save(cache_path)
        return index

    @classmethod
    def combine(cls, indexes: Sequence["EmbeddingIndex"]) -> "EmbeddingIndex":
        """Concatenate multiple embedding indexes into a single searchable index."""
        if not indexes:
            raise ValueError("At least one embedding index must be provided.")

        descriptions: list[str] = []
        codes: list[str] = []
        vocabularies: list[str] = []
        embeddings: list[np.ndarray] = []

        for index in indexes:
            descriptions.extend(index.descriptions)
            codes.extend(index.codes)
            vocabularies.extend(index.vocabularies)
            embeddings.append(np.asarray(index.embeddings))

        return cls(
            model=indexes[0].model,
            descriptions=descriptions,
            codes=codes,
            vocabularies=vocabularies,
            embeddings=(embeddings[0] if len(embeddings) == 1 else np.vstack(embeddings)).astype(np.float32, copy=False)
            if embeddings
            else np.empty((0, 0), dtype=np.float32),
        )

    @classmethod
    def load_for_vocabularies(
        cls,
        codes: pd.DataFrame,
        vocabularies: Sequence[str],
        cache_dir: str | Path,
        model_name: str = DEFAULT_LOCAL_EMBEDDING_MODEL,
        model: EmbeddingEncoder | None = None,
        provider: str = "local",
        dimensions: int | None = None,
        progress_callback: Callable[[int, int], None] | None = None,
        status_callback: Callable[[str], None] | None = None,
    ) -> "EmbeddingIndex":
        """Load per-vocabulary embedding caches and concatenate them.
        Each vocabulary is cached independently by vocabulary, provider,
        model, and output dimensions.
        Selected vocabularies are then combined in memory for retrieval.
        """
        cache_dir = Path(cache_dir)
        indexes: list[EmbeddingIndex] = []
        total = len(vocabularies)
        if progress_callback is not None:
            progress_callback(0, total)
        for position, vocabulary in enumerate(vocabularies, start=1):
            vocab_df = codes[
                codes["vocabulary"].astype(str).str.upper() == vocabulary.upper()
            ].reset_index(drop=True)
            if vocab_df.empty:
                continue
            cache_identity = embedding_cache_key(provider, model_name, dimensions)
            cache_path = cache_dir / f"embeddings_{vocabulary}_{cache_identity}.npz"
            if status_callback is not None:
                status_callback(f"Preparing embeddings for {vocabulary} ({len(vocab_df):,} descriptions)...")
            # Existing release assets contain the legacy local MiniLM indexes.
            # OpenAI indexes are private, paid-to-generate local artifacts and
            # must not be mistaken for downloadable release assets.
            if provider.strip().lower() == "local":
                _ensure_embedding_asset(cache_path)
            if status_callback is not None:
                status_callback(
                    f"Loading and checking saved embeddings for {vocabulary}..."
                    if cache_path.exists()
                    else f"Creating embeddings for {vocabulary} ({len(vocab_df):,} descriptions). "
                    "The first run can take several minutes."
                )
            indexes.append(
                cls.from_cache_or_build(
                    vocab_df,
                    cache_path=cache_path,
                    model_name=model_name,
                    model=model,
                    load_model=False,
                    show_progress_bar=progress_callback is not None,
                    status_callback=(
                        (lambda message, vocabulary=vocabulary: status_callback(f"{vocabulary}: {message}"))
                        if status_callback is not None else None
                    ),
                )
            )
            if progress_callback is not None:
                progress_callback(position, total)
        if status_callback is not None:
            status_callback("Combining terminology embeddings...")
        matrices = [np.asarray(index.embeddings, dtype=np.float32) for index in indexes]
        return cls(
            model=None,
            descriptions=[d for index in indexes for d in index.descriptions],
            codes=[c for index in indexes for c in index.codes],
            vocabularies=[v for index in indexes for v in index.vocabularies],
            embeddings=(matrices[0] if len(matrices) == 1 else np.vstack(matrices))
            if indexes
            else np.empty((0, 0), dtype=np.float32),
        )

    def retrieve(
        self,
        query: str,
        top_k: int | None = None,
        model: EmbeddingEncoder | None = None,
    ) -> list[EmbeddingCandidate]:
        query_model = model or self.model
        if query_model is None:
            raise ValueError("An embedding model is required to embed the query.")
        q_emb = query_model.encode(query, show_progress_bar=False, normalize_embeddings=True)
        scores = self.embeddings @ q_emb  # cosine similarity (both normalised)
        if top_k is not None and len(scores) > top_k:
            top_indices = np.argpartition(scores, -top_k)[-top_k:]
            sorted_local = sorted(top_indices, key=lambda i: -scores[i])[:top_k]
        else:
            sorted_local = np.argsort(-scores)

        return [
            EmbeddingCandidate(
                rank=r + 1,
                code=self.codes[idx],
                description=self.descriptions[idx],
                vocabulary=self.vocabularies[idx],
                score=float(round(scores[idx], 6)),
            )
            for r, idx in enumerate(sorted_local)
        ]

    def score_all(
        self,
        query: str,
        model: EmbeddingEncoder | None = None,
    ) -> list[float]:
        """Return corpus-order cosine scores without ranking or object creation.
        Hybrid retrieval needs every embedding score for global min-max
        normalization, but it does not need the embedding-only ranking. This
        path preserves the same six-decimal scores as :meth:`retrieve` while
        avoiding a full sort and hundreds of thousands of temporary
        ``EmbeddingCandidate`` objects.
        """
        query_model = model or self.model
        if query_model is None:
            raise ValueError("An embedding model is required to embed the query.")
        q_emb = query_model.encode(query, show_progress_bar=False, normalize_embeddings=True)
        scores = self.embeddings @ q_emb
        return [float(round(score, 6)) for score in scores]


def _ordered_subset_positions(
    cached_codes: Sequence[str],
    cached_vocabularies: Sequence[str],
    expected_codes: Sequence[str],
    expected_vocabularies: Sequence[str],
) -> list[int] | None:
    """Locate an order-preserving expected corpus inside a cached corpus."""
    if (
        len(expected_codes) != len(expected_vocabularies)
        or len(expected_codes) > len(cached_codes)
    ):
        return None
    positions: list[int] = []
    expected_index = 0
    for cached_index, (code, vocabulary) in enumerate(
        zip(cached_codes, cached_vocabularies)
    ):
        if expected_index >= len(expected_codes):
            break
        if (
            code == expected_codes[expected_index]
            and vocabulary == expected_vocabularies[expected_index]
        ):
            positions.append(cached_index)
            expected_index += 1
    return positions if expected_index == len(expected_codes) else None


def _release_asset_url(asset_name: str) -> str | None:
    owner = os.getenv("GITHUB_REPOSITORY_OWNER")
    repo = os.getenv("GITHUB_REPOSITORY_NAME")
    tag = os.getenv("GITHUB_RELEASE_TAG")
    if not owner or not repo or not tag:
        return None
    return f"https://github.com/{owner}/{repo}/releases/download/{tag}/{asset_name}"


def _ensure_embedding_asset(cache_path: Path) -> None:
    """Ensure an embedding asset exists locally, downloading it from a GitHub Release if needed."""
    if cache_path.exists():
        return

    asset_name = cache_path.name
    url = _release_asset_url(asset_name)
    if url is None:
        return

    cache_path.parent.mkdir(parents=True, exist_ok=True)
    try:
        with urlopen(url) as response, cache_path.open("wb") as fh:
            fh.write(response.read())
    except URLError:
        if cache_path.exists():
            cache_path.unlink(missing_ok=True)
        raise


def retrieve_embeddings(
    query: str,
    codes: pd.DataFrame,
    model_name: str = DEFAULT_LOCAL_EMBEDDING_MODEL,
    top_k: int | None = None,
) -> list[EmbeddingCandidate]:
    """Build an embedding index on *codes* and retrieve candidates for *query*.
    Parameters
    ----------
    query : str
        Phenotype / EDF description to search.
    codes : pd.DataFrame
        Corpus with a ``description`` column (and optionally ``code``/``vocabulary``).
    model_name : str
        sentence-transformers model identifier.
    top_k : int, optional
        Number of results to return.
    Returns
    -------
    list[EmbeddingCandidate]
    """
    df = codes.copy()
    for col in ("code", "vocabulary"):
        if col not in df.columns:
            df[col] = ""

    descs = df["description"].fillna("").astype(str).tolist()
    model = _create_local_embedding_model(model_name)
    embs = model.encode(descs, show_progress_bar=False, normalize_embeddings=True)
    q_emb = model.encode(query, show_progress_bar=False, normalize_embeddings=True)
    scores = embs @ q_emb

    sorted_idx = np.argsort(-scores)
    if top_k is not None:
        sorted_idx = sorted_idx[:top_k]

    return [
        EmbeddingCandidate(
            rank=r + 1,
            code=df.iloc[idx]["code"],
            description=descs[idx],
            vocabulary=str(df.iloc[idx].get("vocabulary", "")),
            score=float(round(scores[idx], 6)),
        )
        for r, idx in enumerate(sorted_idx)
    ]

	"""Embedding-based retrieval using local or API-backed encoders."""

	from __future__ import annotations

	from collections.abc import Callable
	from dataclasses import dataclass
	import os
	from pathlib import Path
	import re
	from typing import Protocol, Sequence
	from urllib.error import URLError
	from urllib.request import urlopen

	import numpy as np
	import pandas as pd

	DEFAULT_LOCAL_EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
	DEFAULT_OPENAI_EMBEDDING_MODEL = "text-embedding-3-large"


	def _create_local_embedding_model(model_name: str):
	"""Import the optional local stack only when that backend is selected."""
	from sentence_transformers import SentenceTransformer

	return SentenceTransformer(model_name)


	class EmbeddingEncoder(Protocol):
	"""Common interface for local and API-backed embedding models."""

	def encode(self, sentences, **kwargs): ...


	class OpenAIEmbeddingModel:
	"""SentenceTransformer-compatible adapter for OpenAI embeddings."""

	def __init__(
	self,
	model_name: str = DEFAULT_OPENAI_EMBEDDING_MODEL,
	*,
	dimensions: int \| None = 3072,
	batch_size: int = 512,
	) -> None:
	if dimensions is not None and dimensions < 1:
	raise ValueError("Embedding dimensions must be at least 1")
	if batch_size < 1:
	raise ValueError("Embedding batch size must be at least 1")
	self.model_name = model_name
	self.dimensions = dimensions
	self.batch_size = batch_size
	self._client = None

	def _get_client(self):
	if self._client is None:
	import openai

	self._client = openai.OpenAI()
	return self._client

	def encode(
	self,
	sentences,
	*,
	normalize_embeddings: bool = True,
	show_progress_bar: bool = False,
	**_kwargs,
	) -> np.ndarray:
	del show_progress_bar
	single = isinstance(sentences, str)
	inputs = [sentences] if single else [str(item) for item in sentences]
	if not inputs:
	width = self.dimensions or 0
	return np.empty((0, width), dtype=np.float32)

	# Convert each response to float32 immediately. Keeping hundreds of
	# thousands of embeddings as Python float lists can require many times
	# more memory than the final NumPy matrix.
	chunks: list[np.ndarray] = []
	client = self._get_client()
	for start in range(0, len(inputs), self.batch_size):
	request = {
	"model": self.model_name,
	"input": inputs[start : start + self.batch_size],
	"encoding_format": "float",
	}
	if self.dimensions is not None:
	request["dimensions"] = self.dimensions
	response = client.embeddings.create(**request)
	chunks.append(
	np.asarray(
	[
	item.embedding
	for item in sorted(response.data, key=lambda item: item.index)
	],
	dtype=np.float32,
	)
	)

	result = np.vstack(chunks)
	if normalize_embeddings:
	norms = np.linalg.norm(result, axis=1, keepdims=True)
	result = result / np.maximum(norms, 1e-12)
	return result[0] if single else result


	def create_embedding_model(
	provider: str,
	model_name: str,
	*,
	dimensions: int \| None = None,
	batch_size: int = 512,
	) -> EmbeddingEncoder:
	"""Construct a local SentenceTransformer or OpenAI embedding adapter."""
	provider = provider.strip().lower()
	if provider == "local":
	return _create_local_embedding_model(model_name)
	if provider == "openai":
	return OpenAIEmbeddingModel(
	model_name,
	dimensions=dimensions,
	batch_size=batch_size,
	)
	raise ValueError(f"Unsupported embedding provider: {provider!r}")


	def embedding_cache_key(
	provider: str,
	model_name: str,
	dimensions: int \| None = None,
	) -> str:
	"""Return a filesystem-safe cache identity without mixing providers/models."""
	model_key = re.sub(r"[^A-Za-z0-9_.-]+", "_", model_name).strip("_")
	if provider.strip().lower() == "local":
	return model_key
	dimension_key = f"_{dimensions}" if dimensions is not None else ""
	return f"{provider.strip().lower()}_{model_key}{dimension_key}"


	@dataclass(slots=True)
	class EmbeddingCandidate:
	rank: int
	code: str
	description: str
	vocabulary: str
	score: float

	def __lt__(self, other: object) -> bool:
	if not isinstance(other, EmbeddingCandidate):
	return NotImplemented
	return self.score > other.score


	@dataclass
	class EmbeddingIndex:
	"""Wraps embedding model + corpus embeddings for fast cosine-sim search."""

	model: EmbeddingEncoder \| None
	descriptions: list[str]
	codes: list[str]
	vocabularies: list[str]
	embeddings: np.ndarray # shape (n, dim)

	@classmethod
	def from_codes(
	cls,
	codes: pd.DataFrame,
	model_name: str = DEFAULT_LOCAL_EMBEDDING_MODEL,
	model: EmbeddingEncoder \| None = None,
	show_progress_bar: bool = False,
	) -> "EmbeddingIndex":
	df = codes.copy()
	for col in ("code", "vocabulary"):
	if col not in df.columns:
	df[col] = ""
	embedding_model = model or _create_local_embedding_model(model_name)
	descs = df["description"].fillna("").astype(str).tolist()
	embs = embedding_model.encode(
	descs,
	show_progress_bar=show_progress_bar,
	normalize_embeddings=True,
	)
	return cls(
	model=embedding_model,
	descriptions=descs,
	codes=df["code"].astype(str).tolist(),
	vocabularies=df["vocabulary"].astype(str).tolist(),
	embeddings=np.asarray(embs, dtype=np.float32),
	)

	def save(self, path: str \| Path) -> None:
	"""Persist the corpus embeddings + metadata to disk (npz) so they can
	be reloaded without re-running the embedding model."""
	path = Path(path)
	path.parent.mkdir(parents=True, exist_ok=True)
	np.savez_compressed(
	path,
	embeddings=np.asarray(self.embeddings, dtype=np.float32),
	codes=np.array(self.codes, dtype=object),
	descriptions=np.array(self.descriptions, dtype=object),
	vocabularies=np.array(self.vocabularies, dtype=object),
	)

	@classmethod
	def load(
	cls,
	path: str \| Path,
	model_name: str = DEFAULT_LOCAL_EMBEDDING_MODEL,
	model: EmbeddingEncoder \| None = None,
	load_model: bool = True,
	) -> "EmbeddingIndex":
	"""Load a previously-saved embedding index from disk."""
	with np.load(path, allow_pickle=True) as data:
	return cls(
	model=model or (_create_local_embedding_model(model_name) if load_model else None),
	descriptions=list(data["descriptions"]),
	codes=list(data["codes"]),
	vocabularies=list(data["vocabularies"]),
	embeddings=np.asarray(data["embeddings"], dtype=np.float32),
	)

	@classmethod
	def from_cache_or_build(
	cls,
	codes: pd.DataFrame,
	cache_path: str \| Path,
	model_name: str = DEFAULT_LOCAL_EMBEDDING_MODEL,
	model: EmbeddingEncoder \| None = None,
	load_model: bool = True,
	show_progress_bar: bool = False,
	status_callback: Callable[[str], None] \| None = None,
	) -> "EmbeddingIndex":
	"""Load a cached embedding index if present, otherwise build it from
	codes and cache the result for subsequent runs.

	This is the key mechanism that avoids re-embedding large reference
	code systems (e.g. the ~74k-row full ICD-10-CM corpus) on every run.
	"""
	cache_path = Path(cache_path)
	if cache_path.exists():
	cached = cls.load(
	cache_path,
	model_name=model_name,
	model=model,
	load_model=load_model,
	)
	expected_codes = codes.get("code", pd.Series(dtype=str)).astype(str).tolist()
	expected_vocabularies = (
	codes.get("vocabulary", pd.Series(dtype=str)).astype(str).tolist()
	)
	expected_descriptions = codes["description"].fillna("").astype(str).tolist()
	if (
	cached.codes == expected_codes
	and cached.vocabularies == expected_vocabularies
	and cached.descriptions == expected_descriptions
	):
	return cached
	subset_positions = _ordered_subset_positions(
	cached.codes,
	cached.vocabularies,
	expected_codes,
	expected_vocabularies,
	)
	if subset_positions is not None and all(
	cached.descriptions[position] == description
	for position, description in zip(subset_positions, expected_descriptions)
	):
	# Corpus preprocessing can intentionally remove non-codable
	# dictionary rows (for example ICD range headers). Reuse the
	# matching cached vectors instead of re-embedding every code.
	subset = cls(
	model=cached.model,
	descriptions=expected_descriptions,
	codes=expected_codes,
	vocabularies=expected_vocabularies,
	embeddings=np.asarray(cached.embeddings[subset_positions], dtype=np.float32),
	)
	subset.save(cache_path)
	return subset
	if status_callback is not None:
	status_callback(
	f"Creating embeddings for {len(codes):,} descriptions; no matching saved index. "
	"The first run can take several minutes."
	)
	index = cls.from_codes(
	codes,
	model_name=model_name,
	model=model,
	show_progress_bar=show_progress_bar,
	)
	index.save(cache_path)
	return index

	@classmethod
	def combine(cls, indexes: Sequence["EmbeddingIndex"]) -> "EmbeddingIndex":
	"""Concatenate multiple embedding indexes into a single searchable index."""
	if not indexes:
	raise ValueError("At least one embedding index must be provided.")

	descriptions: list[str] = []
	codes: list[str] = []
	vocabularies: list[str] = []
	embeddings: list[np.ndarray] = []

	for index in indexes:
	descriptions.extend(index.descriptions)
	codes.extend(index.codes)
	vocabularies.extend(index.vocabularies)
	embeddings.append(np.asarray(index.embeddings))

	return cls(
	model=indexes[0].model,
	descriptions=descriptions,
	codes=codes,
	vocabularies=vocabularies,
	embeddings=(embeddings[0] if len(embeddings) == 1 else np.vstack(embeddings)).astype(np.float32, copy=False)
	if embeddings
	else np.empty((0, 0), dtype=np.float32),
	)

	@classmethod
	def load_for_vocabularies(
	cls,
	codes: pd.DataFrame,
	vocabularies: Sequence[str],
	cache_dir: str \| Path,
	model_name: str = DEFAULT_LOCAL_EMBEDDING_MODEL,
	model: EmbeddingEncoder \| None = None,
	provider: str = "local",
	dimensions: int \| None = None,
	progress_callback: Callable[[int, int], None] \| None = None,
	status_callback: Callable[[str], None] \| None = None,
	) -> "EmbeddingIndex":
	"""Load per-vocabulary embedding caches and concatenate them.

	Each vocabulary is cached independently by vocabulary, provider,
	model, and output dimensions.
	Selected vocabularies are then combined in memory for retrieval.
	"""
	cache_dir = Path(cache_dir)
	indexes: list[EmbeddingIndex] = []
	total = len(vocabularies)
	if progress_callback is not None:
	progress_callback(0, total)
	for position, vocabulary in enumerate(vocabularies, start=1):
	vocab_df = codes[
	codes["vocabulary"].astype(str).str.upper() == vocabulary.upper()
	].reset_index(drop=True)
	if vocab_df.empty:
	continue
	cache_identity = embedding_cache_key(provider, model_name, dimensions)
	cache_path = cache_dir / f"embeddings_{vocabulary}_{cache_identity}.npz"
	if status_callback is not None:
	status_callback(f"Preparing embeddings for {vocabulary} ({len(vocab_df):,} descriptions)...")
	# Existing release assets contain the legacy local MiniLM indexes.
	# OpenAI indexes are private, paid-to-generate local artifacts and
	# must not be mistaken for downloadable release assets.
	if provider.strip().lower() == "local":
	_ensure_embedding_asset(cache_path)
	if status_callback is not None:
	status_callback(
	f"Loading and checking saved embeddings for {vocabulary}..."
	if cache_path.exists()
	else f"Creating embeddings for {vocabulary} ({len(vocab_df):,} descriptions). "
	"The first run can take several minutes."
	)
	indexes.append(
	cls.from_cache_or_build(
	vocab_df,
	cache_path=cache_path,
	model_name=model_name,
	model=model,
	load_model=False,
	show_progress_bar=progress_callback is not None,
	status_callback=(
	(lambda message, vocabulary=vocabulary: status_callback(f"{vocabulary}: {message}"))
	if status_callback is not None else None
	),
	)
	)
	if progress_callback is not None:
	progress_callback(position, total)
	if status_callback is not None:
	status_callback("Combining terminology embeddings...")
	matrices = [np.asarray(index.embeddings, dtype=np.float32) for index in indexes]
	return cls(
	model=None,
	descriptions=[d for index in indexes for d in index.descriptions],
	codes=[c for index in indexes for c in index.codes],
	vocabularies=[v for index in indexes for v in index.vocabularies],
	embeddings=(matrices[0] if len(matrices) == 1 else np.vstack(matrices))
	if indexes
	else np.empty((0, 0), dtype=np.float32),
	)

	def retrieve(
	self,
	query: str,
	top_k: int \| None = None,
	model: EmbeddingEncoder \| None = None,
	) -> list[EmbeddingCandidate]:
	query_model = model or self.model
	if query_model is None:
	raise ValueError("An embedding model is required to embed the query.")
	q_emb = query_model.encode(query, show_progress_bar=False, normalize_embeddings=True)
	scores = self.embeddings @ q_emb # cosine similarity (both normalised)
	if top_k is not None and len(scores) > top_k:
	top_indices = np.argpartition(scores, -top_k)[-top_k:]
	sorted_local = sorted(top_indices, key=lambda i: -scores[i])[:top_k]
	else:
	sorted_local = np.argsort(-scores)

	return [
	EmbeddingCandidate(
	rank=r + 1,
	code=self.codes[idx],
	description=self.descriptions[idx],
	vocabulary=self.vocabularies[idx],
	score=float(round(scores[idx], 6)),
	)
	for r, idx in enumerate(sorted_local)
	]

	def score_all(
	self,
	query: str,
	model: EmbeddingEncoder \| None = None,
	) -> list[float]:
	"""Return corpus-order cosine scores without ranking or object creation.

	Hybrid retrieval needs every embedding score for global min-max
	normalization, but it does not need the embedding-only ranking. This
	path preserves the same six-decimal scores as :meth:`retrieve` while
	avoiding a full sort and hundreds of thousands of temporary
	``EmbeddingCandidate`` objects.
	"""
	query_model = model or self.model
	if query_model is None:
	raise ValueError("An embedding model is required to embed the query.")
	q_emb = query_model.encode(query, show_progress_bar=False, normalize_embeddings=True)
	scores = self.embeddings @ q_emb
	return [float(round(score, 6)) for score in scores]


	def _ordered_subset_positions(
	cached_codes: Sequence[str],
	cached_vocabularies: Sequence[str],
	expected_codes: Sequence[str],
	expected_vocabularies: Sequence[str],
	) -> list[int] \| None:
	"""Locate an order-preserving expected corpus inside a cached corpus."""
	if (
	len(expected_codes) != len(expected_vocabularies)
	or len(expected_codes) > len(cached_codes)
	):
	return None
	positions: list[int] = []
	expected_index = 0
	for cached_index, (code, vocabulary) in enumerate(
	zip(cached_codes, cached_vocabularies)
	):
	if expected_index >= len(expected_codes):
	break
	if (
	code == expected_codes[expected_index]
	and vocabulary == expected_vocabularies[expected_index]
	):
	positions.append(cached_index)
	expected_index += 1
	return positions if expected_index == len(expected_codes) else None


	def _release_asset_url(asset_name: str) -> str \| None:
	owner = os.getenv("GITHUB_REPOSITORY_OWNER")
	repo = os.getenv("GITHUB_REPOSITORY_NAME")
	tag = os.getenv("GITHUB_RELEASE_TAG")
	if not owner or not repo or not tag:
	return None
	return f"https://github.com/{owner}/{repo}/releases/download/{tag}/{asset_name}"


	def _ensure_embedding_asset(cache_path: Path) -> None:
	"""Ensure an embedding asset exists locally, downloading it from a GitHub Release if needed."""
	if cache_path.exists():
	return

	asset_name = cache_path.name
	url = _release_asset_url(asset_name)
	if url is None:
	return

	cache_path.parent.mkdir(parents=True, exist_ok=True)
	try:
	with urlopen(url) as response, cache_path.open("wb") as fh:
	fh.write(response.read())
	except URLError:
	if cache_path.exists():
	cache_path.unlink(missing_ok=True)
	raise


	def retrieve_embeddings(
	query: str,
	codes: pd.DataFrame,
	model_name: str = DEFAULT_LOCAL_EMBEDDING_MODEL,
	top_k: int \| None = None,
	) -> list[EmbeddingCandidate]:
	"""Build an embedding index on codes and retrieve candidates for query.

	Parameters
	----------
	query : str
	Phenotype / EDF description to search.
	codes : pd.DataFrame
	Corpus with a ``description`` column (and optionally ``code``/``vocabulary``).
	model_name : str
	sentence-transformers model identifier.
	top_k : int, optional
	Number of results to return.

	Returns
	-------
	list[EmbeddingCandidate]
	"""
	df = codes.copy()
	for col in ("code", "vocabulary"):
	if col not in df.columns:
	df[col] = ""

	descs = df["description"].fillna("").astype(str).tolist()
	model = _create_local_embedding_model(model_name)
	embs = model.encode(descs, show_progress_bar=False, normalize_embeddings=True)
	q_emb = model.encode(query, show_progress_bar=False, normalize_embeddings=True)
	scores = embs @ q_emb

	sorted_idx = np.argsort(-scores)
	if top_k is not None:
	sorted_idx = sorted_idx[:top_k]

	return [
	EmbeddingCandidate(
	rank=r + 1,
	code=df.iloc[idx]["code"],
	description=descs[idx],
	vocabulary=str(df.iloc[idx].get("vocabulary", "")),
	score=float(round(scores[idx], 6)),
	)
	for r, idx in enumerate(sorted_idx)
	]