Spaces:
Sleeping
Sleeping
File size: 4,428 Bytes
c2b7494 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 | """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]
@dataclass(slots=True)
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
|