Spaces:
Sleeping
Sleeping
| """Hybrid retrieval combining lexical (BM25) and embedding (semantic) scores.""" | |
| from __future__ import annotations | |
| from dataclasses import dataclass | |
| import pandas as pd | |
| from src.retrieval.embeddings import EmbeddingEncoder, EmbeddingIndex, retrieve_embeddings | |
| from src.retrieval.lexical import build_lexical_index, score_lexical | |
| class HybridCandidate: | |
| rank: int | |
| code: str | |
| description: str | |
| vocabulary: str | |
| lexical_score: float | |
| embedding_score: float | |
| score: float | |
| def __lt__(self, other: object) -> bool: | |
| if not isinstance(other, HybridCandidate): | |
| return NotImplemented | |
| return self.score > other.score | |
| def _min_max_normalize(values: list[float]) -> list[float]: | |
| if not values: | |
| return values | |
| lo, hi = min(values), max(values) | |
| if hi - lo < 1e-12: | |
| return [0.0 for _ in values] | |
| return [(v - lo) / (hi - lo) for v in values] | |
| def hybrid_retrieval( | |
| query: str, | |
| codes: pd.DataFrame, | |
| lexical_weight: float = 0.5, | |
| embedding_weight: float = 0.5, | |
| embedding_model: str = "sentence-transformers/all-MiniLM-L6-v2", | |
| top_k: int | None = 25, | |
| per_vocabulary_top_k: bool = True, | |
| embedding_index: EmbeddingIndex | None = None, | |
| query_model: EmbeddingEncoder | None = None, | |
| lexical_query: str | None = None, | |
| embedding_query: str | None = None, | |
| ) -> list[HybridCandidate]: | |
| """Combine BM25 lexical retrieval with embedding-based semantic retrieval. | |
| Both signals are min-max normalized to [0, 1] and combined with the given | |
| weights before ranking. The combined ranking is truncated to *top_k*. | |
| Parameters | |
| ---------- | |
| query : str | |
| Phenotype / EDF text used as the search query. | |
| codes : pd.DataFrame | |
| Candidate code corpus with ``code``, ``description``, ``vocabulary`` columns. | |
| lexical_weight, embedding_weight : float | |
| Relative weighting of each retrieval signal. Need not sum to 1. | |
| embedding_model : str | |
| sentence-transformers model identifier used when *embedding_index* is None. | |
| top_k : int, optional | |
| Number of ranked results to return per vocabulary. ``None`` returns the | |
| full corpus. | |
| per_vocabulary_top_k : bool | |
| If ``True`` (default), apply ``top_k`` separately within each selected | |
| vocabulary. If ``False``, apply ``top_k`` after combining all selected | |
| vocabularies. | |
| embedding_index : EmbeddingIndex, optional | |
| Pre-built embedding index to reuse across multiple queries. | |
| query_model : EmbeddingEncoder, optional | |
| Cached local or API-backed encoder used to embed the EDF/query text. | |
| lexical_query : str, optional | |
| Text used for the BM25 lexical score. If omitted, *query* is used. | |
| embedding_query : str, optional | |
| Text used for the embedding score. If omitted, *query* is used. | |
| Returns | |
| ------- | |
| list[HybridCandidate] | |
| """ | |
| if codes.empty: | |
| return [] | |
| df = codes.reset_index(drop=True).copy() | |
| for col in ("code", "vocabulary"): | |
| if col not in df.columns: | |
| df[col] = "" | |
| lexical_query = lexical_query or query | |
| embedding_query = embedding_query or query | |
| embedding_scores_by_code: dict[tuple[str, str], float] | None = None | |
| if embedding_index is not None: | |
| # Score the combined corpus exactly once. Previously this full matrix | |
| # multiplication, full sort, and candidate-object construction ran | |
| # once per vocabulary even though every pass produced identical raw | |
| # scores. The corpus is deduplicated on (code, vocabulary), so the map | |
| # is identical to the maps produced by the former repeated searches. | |
| all_embedding_scores = embedding_index.score_all(embedding_query, model=query_model) | |
| embedding_scores_by_code = { | |
| (code, vocabulary): score | |
| for code, vocabulary, score in zip( | |
| embedding_index.codes, | |
| embedding_index.vocabularies, | |
| all_embedding_scores, | |
| ) | |
| } | |
| per_vocab_candidates: list[ | |
| tuple[str, list[tuple[str, str, str, float, float]]] | |
| ] = [] | |
| for vocabulary, vocab_df in df.groupby("vocabulary", sort=False): | |
| lexical_index = build_lexical_index(vocab_df) | |
| lexical_scores = score_lexical(lexical_query, lexical_index, len(vocab_df)) | |
| if embedding_scores_by_code is None: | |
| embedding_results = retrieve_embeddings(embedding_query, vocab_df, model_name=embedding_model, top_k=None) | |
| vocabulary_embedding_scores = { | |
| (c.code, c.vocabulary): c.score for c in embedding_results | |
| } | |
| else: | |
| vocabulary_embedding_scores = embedding_scores_by_code | |
| codes_values = vocab_df["code"].tolist() | |
| description_values = vocab_df["description"].tolist() | |
| vocabulary_values = vocab_df["vocabulary"].tolist() | |
| vocab_candidates = [ | |
| ( | |
| code, | |
| description, | |
| candidate_vocabulary, | |
| lexical_score, | |
| vocabulary_embedding_scores.get((code, candidate_vocabulary), 0.0), | |
| ) | |
| for code, description, candidate_vocabulary, lexical_score in zip( | |
| codes_values, | |
| description_values, | |
| vocabulary_values, | |
| lexical_scores, | |
| ) | |
| ] | |
| per_vocab_candidates.append((vocabulary, vocab_candidates)) | |
| all_lexical = [ | |
| lexical | |
| for _, vocab_candidates in per_vocab_candidates | |
| for _, _, _, lexical, _ in vocab_candidates | |
| ] | |
| all_embedding = [ | |
| embedding | |
| for _, vocab_candidates in per_vocab_candidates | |
| for _, _, _, _, embedding in vocab_candidates | |
| ] | |
| norm_all_lexical = _min_max_normalize(all_lexical) | |
| norm_all_embedding = _min_max_normalize(all_embedding) | |
| # Keep lightweight tuples until after top-k selection. Constructing a | |
| # dataclass (and formerly a Pandas Series) for all 620k rows dominated the | |
| # runtime even when only a handful of results were requested. | |
| selected: list[tuple[str, str, str, float, float, float]] = [] | |
| flat_idx = 0 | |
| for _, vocab_candidates in per_vocab_candidates: | |
| scored: list[tuple[str, str, str, float, float, float]] = [] | |
| for idx, (code, description, vocabulary, _, _) in enumerate(vocab_candidates): | |
| norm_lexical = norm_all_lexical[flat_idx + idx] | |
| norm_embedding = norm_all_embedding[flat_idx + idx] | |
| score = lexical_weight * norm_lexical + embedding_weight * norm_embedding | |
| scored.append( | |
| ( | |
| code, | |
| description, | |
| vocabulary, | |
| round(norm_lexical, 6), | |
| round(norm_embedding, 6), | |
| round(score, 6), | |
| ) | |
| ) | |
| if per_vocabulary_top_k and top_k is not None: | |
| scored.sort(key=lambda candidate: -candidate[5]) | |
| scored = scored[:top_k] | |
| selected.extend(scored) | |
| flat_idx += len(vocab_candidates) | |
| selected.sort(key=lambda candidate: -candidate[5]) | |
| if not per_vocabulary_top_k and top_k is not None: | |
| selected = selected[:top_k] | |
| return [ | |
| HybridCandidate( | |
| rank=rank, | |
| code=code, | |
| description=description, | |
| vocabulary=vocabulary, | |
| lexical_score=lexical_score, | |
| embedding_score=embedding_score, | |
| score=score, | |
| ) | |
| for rank, (code, description, vocabulary, lexical_score, embedding_score, score) | |
| in enumerate(selected, start=1) | |
| ] | |