87 lines
2.9 KiB
Python
87 lines
2.9 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
from pathlib import Path
|
|
|
|
from .models import ParentDocument, RetrievalDocument, SourceRef
|
|
|
|
|
|
def _read_jsonl(path: Path) -> list[dict]:
|
|
with path.open(encoding="utf-8") as handle:
|
|
return [json.loads(line) for line in handle if line.strip()]
|
|
|
|
|
|
def _source_ref(raw: dict) -> SourceRef:
|
|
bbox = raw.get("bbox")
|
|
page_range = raw.get("page_range")
|
|
printed_page_range = raw.get("printed_page_range")
|
|
return SourceRef(
|
|
physical_page=int(raw["physical_page"]),
|
|
precision=raw["precision"],
|
|
block_id=raw.get("block_id"),
|
|
bbox=tuple(bbox) if bbox else None,
|
|
source_crop=raw.get("source_crop"),
|
|
page_range=tuple(page_range) if page_range else None,
|
|
printed_page=raw.get("printed_page"),
|
|
printed_page_range=(
|
|
tuple(printed_page_range) if printed_page_range else None
|
|
),
|
|
)
|
|
|
|
|
|
def load_documents(path: Path) -> list[RetrievalDocument]:
|
|
documents = []
|
|
for raw in _read_jsonl(path):
|
|
documents.append(RetrievalDocument(
|
|
doc_id=raw["doc_id"],
|
|
drug_id=raw["drug_id"],
|
|
kind=raw["kind"],
|
|
text=raw["text"],
|
|
section_key=raw["section_key"],
|
|
source_refs=tuple(_source_ref(item) for item in raw.get("source_refs", [])),
|
|
parent_id=raw.get("parent_id"),
|
|
requires_visual_check=raw.get("requires_visual_check", False),
|
|
drug_name=raw.get("drug_name"),
|
|
))
|
|
return documents
|
|
|
|
|
|
def load_parents(path: Path) -> list[ParentDocument]:
|
|
parents = []
|
|
for raw in _read_jsonl(path):
|
|
parents.append(ParentDocument(
|
|
parent_id=raw["logical_table_id"],
|
|
kind=raw["kind"],
|
|
text=raw["markdown"],
|
|
source_refs=tuple(_source_ref(item) for item in raw.get("source_refs", [])),
|
|
requires_visual_check=raw.get("requires_visual_check", False),
|
|
))
|
|
return parents
|
|
|
|
|
|
def load_aliases(path: Path | None) -> dict[str, set[str]]:
|
|
if path is None:
|
|
return {}
|
|
raw = json.loads(path.read_text(encoding="utf-8"))
|
|
if isinstance(raw, dict) and "entities" in raw:
|
|
return {
|
|
entity["drug_id"]: set(entity["aliases"])
|
|
for entity in raw["entities"]
|
|
}
|
|
return {drug_id: set(aliases) for drug_id, aliases in raw.items()}
|
|
|
|
|
|
def build_drug_catalog(
|
|
documents: list[RetrievalDocument],
|
|
extra_aliases: dict[str, set[str]] | None = None,
|
|
) -> dict[str, set[str]]:
|
|
catalog: dict[str, set[str]] = {}
|
|
for document in documents:
|
|
aliases = catalog.setdefault(document.drug_id, set())
|
|
aliases.add(document.drug_id.replace("_", " "))
|
|
if document.drug_name:
|
|
aliases.add(document.drug_name)
|
|
for drug_id, aliases in (extra_aliases or {}).items():
|
|
catalog.setdefault(drug_id, set()).update(aliases)
|
|
return catalog
|