98 lines
3.5 KiB
Python
98 lines
3.5 KiB
Python
from __future__ import annotations
|
|
|
|
import re
|
|
import unicodedata
|
|
from collections import Counter
|
|
from math import log
|
|
|
|
from .models import ParentDocument, RetrievalDocument, SearchHit
|
|
|
|
WORD_RE = re.compile(r"\w+", re.UNICODE)
|
|
|
|
|
|
def _normalized(text: str) -> str:
|
|
return " ".join(WORD_RE.findall(unicodedata.normalize("NFKC", text).casefold()))
|
|
|
|
|
|
def _terms(text: str) -> set[str]:
|
|
return set(_normalized(text).split())
|
|
|
|
|
|
def _char_ngrams(text: str, size: int = 3) -> set[str]:
|
|
normalized = _normalized(text)
|
|
if len(normalized) <= size:
|
|
return {normalized} if normalized else set()
|
|
return {
|
|
normalized[index:index + size]
|
|
for index in range(len(normalized) - size + 1)
|
|
}
|
|
|
|
|
|
class InMemoryLexicalRetriever:
|
|
"""Deterministic test/fallback retriever, not the production neural backend."""
|
|
|
|
def __init__(self, documents: list[RetrievalDocument]) -> None:
|
|
self._documents = tuple(documents)
|
|
self._term_counts = {
|
|
document.doc_id: Counter(_normalized(document.text).split())
|
|
for document in self._documents
|
|
}
|
|
self._average_length = (
|
|
sum(sum(counts.values()) for counts in self._term_counts.values())
|
|
/ max(1, len(self._term_counts))
|
|
)
|
|
document_frequency: Counter[str] = Counter()
|
|
for counts in self._term_counts.values():
|
|
document_frequency.update(counts.keys())
|
|
self._document_frequency = document_frequency
|
|
|
|
def search(self, query: str, drug_id: str, limit: int) -> list[SearchHit]:
|
|
query_terms = _terms(query)
|
|
query_ngrams = _char_ngrams(query)
|
|
candidates = []
|
|
for document in self._documents:
|
|
if document.drug_id != drug_id:
|
|
continue
|
|
counts = self._term_counts[document.doc_id]
|
|
bm25 = self._bm25(query_terms, counts)
|
|
ngrams = _char_ngrams(document.text)
|
|
char_score = len(query_ngrams & ngrams) / max(1, len(query_ngrams))
|
|
if bm25 > 0 or char_score > 0:
|
|
candidates.append((document, bm25, char_score))
|
|
max_bm25 = max((row[1] for row in candidates), default=0.0)
|
|
hits = [
|
|
SearchHit(
|
|
document=document,
|
|
score=0.8 * (bm25 / max_bm25 if max_bm25 else 0.0) + 0.2 * char_score,
|
|
)
|
|
for document, bm25, char_score in candidates
|
|
]
|
|
return sorted(hits, key=lambda hit: (-hit.score, hit.document.doc_id))[:limit]
|
|
|
|
def _bm25(self, query_terms: set[str], counts: Counter[str]) -> float:
|
|
total_documents = len(self._documents)
|
|
document_length = sum(counts.values())
|
|
score = 0.0
|
|
for term in query_terms:
|
|
frequency = counts.get(term, 0)
|
|
if not frequency:
|
|
continue
|
|
document_frequency = self._document_frequency[term]
|
|
inverse_frequency = log(
|
|
1 + (total_documents - document_frequency + 0.5)
|
|
/ (document_frequency + 0.5)
|
|
)
|
|
denominator = frequency + 1.5 * (
|
|
1 - 0.75 + 0.75 * document_length / max(1.0, self._average_length)
|
|
)
|
|
score += inverse_frequency * frequency * 2.5 / denominator
|
|
return score
|
|
|
|
|
|
class InMemoryParentStore:
|
|
def __init__(self, parents: list[ParentDocument]) -> None:
|
|
self._parents = {parent.parent_id: parent for parent in parents}
|
|
|
|
def get(self, parent_id: str) -> ParentDocument | None:
|
|
return self._parents.get(parent_id)
|