63 lines
2.2 KiB
Python
63 lines
2.2 KiB
Python
"""Bounded sibling expansion for split clinical sections."""
|
|
from __future__ import annotations
|
|
|
|
from collections.abc import Sequence
|
|
|
|
from .models import RetrievalDocument, SearchHit
|
|
|
|
|
|
def _compatible(left: RetrievalDocument, right: RetrievalDocument) -> bool:
|
|
"""Prevent expansion across two explicitly different clinical scopes."""
|
|
if left.drug_id != right.drug_id or left.section_key != right.section_key:
|
|
return False
|
|
if left.context_labels and right.context_labels:
|
|
return bool(set(left.context_labels) & set(right.context_labels))
|
|
return not left.context_labels and not right.context_labels
|
|
|
|
|
|
def expand_siblings(
|
|
seeds: Sequence[SearchHit],
|
|
documents: Sequence[RetrievalDocument],
|
|
*,
|
|
window: int = 1,
|
|
limit: int | None = None,
|
|
) -> list[SearchHit]:
|
|
"""Add adjacent parts without crossing drug, section, or scope labels.
|
|
|
|
Each seed remains first in its group. Neighbours retain the seed score so
|
|
later context packing treats them as context for that match, not as a new
|
|
independently scored retrieval result.
|
|
"""
|
|
if window < 0:
|
|
raise ValueError("window must be non-negative")
|
|
if limit is not None and limit < 0:
|
|
raise ValueError("limit must be non-negative or None")
|
|
|
|
by_position = {
|
|
(doc.drug_id, doc.section_key, doc.part_index): doc
|
|
for doc in documents
|
|
if doc.part_index is not None
|
|
}
|
|
output: list[SearchHit] = []
|
|
seen: set[str] = set()
|
|
|
|
for seed in seeds:
|
|
doc = seed.document
|
|
group = [doc]
|
|
if doc.part_index is not None and window:
|
|
group = []
|
|
for index in range(doc.part_index - window, doc.part_index + window + 1):
|
|
sibling = by_position.get((doc.drug_id, doc.section_key, index))
|
|
if sibling is not None and _compatible(doc, sibling):
|
|
group.append(sibling)
|
|
group.sort(key=lambda item: item.part_index if item.part_index is not None else 0)
|
|
|
|
for item in group:
|
|
if item.doc_id in seen:
|
|
continue
|
|
seen.add(item.doc_id)
|
|
output.append(SearchHit(item, seed.score))
|
|
if limit is not None and len(output) >= limit:
|
|
return output
|
|
return output
|