Spaces:
Sleeping
Sleeping
| """BM25 lexical retrieval over code descriptions.""" | |
| from __future__ import annotations | |
| import re | |
| from dataclasses import dataclass | |
| import numpy as np | |
| import pandas as pd | |
| from rank_bm25 import BM25Okapi | |
| def _tokenize(text: str) -> list[str]: | |
| """Simple whitespace + punctuation tokenizer for BM25 indexing.""" | |
| if not text or not isinstance(text, str): | |
| return [] | |
| text = re.sub(r"[^\w\s]", " ", text.lower()) | |
| return [tok for tok in text.split() if len(tok) > 1] | |
| class LexicalCandidate: | |
| rank: int | |
| code: str | |
| description: str | |
| vocabulary: str | |
| score: float | |
| def __lt__(self, other: object) -> bool: | |
| if not isinstance(other, LexicalCandidate): | |
| return NotImplemented | |
| return self.score > other.score # higher = more relevant | |
| def build_lexical_index(codes: pd.DataFrame) -> BM25Okapi: | |
| """Build a BM25 index from the code corpus descriptions. | |
| Parameters | |
| ---------- | |
| codes : pd.DataFrame | |
| Must contain ``description`` column with full description strings. | |
| Returns | |
| ------- | |
| BM25Okapi | |
| The built tokenised index. | |
| """ | |
| doc_tokens: list[list[str]] = [] | |
| for desc in codes["description"].astype(str): | |
| tokens = _tokenize(desc) | |
| doc_tokens.append(tokens if tokens else [desc]) | |
| return BM25Okapi(doc_tokens) | |
| def score_lexical(query: str, index: BM25Okapi, corpus_size: int) -> list[float]: | |
| """Return corpus-order BM25 scores without allocating result objects.""" | |
| query_tokens = _tokenize(query) | |
| if not query_tokens: | |
| return [0.0] * corpus_size | |
| return [round(float(score), 6) for score in index.get_scores(query_tokens)] | |
| def retrieve_lexical( | |
| query: str, | |
| codes: pd.DataFrame | list[dict], | |
| index: BM25Okapi | None = None, | |
| top_k: int | None = None, | |
| ) -> list[LexicalCandidate]: | |
| """Retrieve candidate codes ranked by BM25 lexical similarity. | |
| Parameters | |
| ---------- | |
| query : str | |
| The phenotype description / EDF summary to search against. | |
| codes : pd.DataFrame or list[dict] | |
| Code corpus containing ``description`` (mandatory) and optionally | |
| ``code``, ``vocabulary`` columns (DataFrame) or dicts with those keys. | |
| index : BM25Okapi, optional | |
| Pre-built index. If omitted one is built on-the-fly from *codes*. | |
| top_k : int, optional | |
| Maximum number of results to return. Defaults to the full corpus size. | |
| Returns | |
| ------- | |
| list[LexicalCandidate] | |
| Ranked candidates sorted descending by BM25 score. | |
| """ | |
| if isinstance(codes, pd.DataFrame): | |
| df = codes.copy() | |
| keys = ("description", "code", "vocabulary") | |
| for k in keys[1:]: | |
| if k not in df.columns: | |
| df[k] = "" # type: ignore[operator] | |
| else: | |
| df = pd.DataFrame(codes) | |
| for k in ("code", "vocabulary"): | |
| if k not in df.columns: | |
| df[k] = "" | |
| if index is None: | |
| index = build_lexical_index(df) | |
| query_tokens = _tokenize(query) | |
| if not query_tokens: | |
| # Fallback: rank all codes uniformly (no lexical signal) | |
| return [ | |
| LexicalCandidate( | |
| rank=0, code=r["code"], description=r["description"], vocabulary=r["vocabulary"], score=0.0 | |
| ) | |
| for i, r in df.iterrows() | |
| ] | |
| scores = index.get_scores(query_tokens).tolist() # type: ignore[union-attr] | |
| if top_k is not None: | |
| top_indices = np.argpartition(scores, -top_k)[-top_k:] | |
| top_scores = [scores[i] for i in sorted(top_indices, key=lambda x: -scores[x])] | |
| results: list[LexicalCandidate] = [ | |
| LexicalCandidate( | |
| rank=r + 1, | |
| code=df.iloc[idx]["code"], | |
| description=df.iloc[idx]["description"], | |
| vocabulary=df.iloc[idx].get("vocabulary", ""), | |
| score=round(float(s), 6), | |
| ) | |
| for r, (idx, s) in enumerate(sorted(enumerate(top_scores), key=lambda x: -x[1])) | |
| ] | |
| else: | |
| results = [ | |
| LexicalCandidate( | |
| rank=r + 1, | |
| code=df.iloc[r]["code"], | |
| description=df.iloc[r]["description"], | |
| vocabulary=df.iloc[r].get("vocabulary", ""), | |
| score=round(float(scores[r]), 6), | |
| ) | |
| for r in range(len(df)) | |
| ] | |
| return results | |