544 lines
22 KiB
Python
544 lines
22 KiB
Python
"""Section -> chunk logic (pure; no filesystem, no embedding client).
|
|
|
|
ADR 0004: chunk unit is `(drug_id, section_key)`. A section under the token
|
|
ceiling becomes one chunk verbatim. Only the long-tail sections above it are
|
|
sub-chunked, with a sentence-boundary-aware sliding window.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import re
|
|
from dataclasses import dataclass
|
|
from typing import Dict, Iterable, Iterator, List, Mapping, Optional, Sequence
|
|
|
|
from ..segment.models import Monograph, SectionSpan, TableBlock
|
|
from ..tables.classify import SHAPE_FORMULA_2D
|
|
from .models import (
|
|
CHUNK_KIND_BLOCK_DESCRIPTOR,
|
|
CHUNK_KIND_PROSE,
|
|
Chunk,
|
|
ChunkAttachment,
|
|
)
|
|
from .sentences import split_sentences
|
|
from .tokens import TokenCounter, count_tokens
|
|
|
|
CEILING_TOKENS = 800
|
|
TARGET_TOKENS = 650
|
|
OVERLAP_TOKENS = 65
|
|
|
|
KIND_TABLE = "table"
|
|
KIND_FORMULA = "formula"
|
|
|
|
# A header row is only safe to embed when it is genuinely a row of labels.
|
|
# Measured on the corpus: 42 of 124 simple-table headers (34%) contain a
|
|
# digit, and AMIODARON's (physical page 183) is
|
|
# "Thời gian liệu pháp tĩnh mạch Liều 720 mg/ngày (0,5 mg/phút)" — a dose,
|
|
# inside what pdfplumber called a header, from an extraction never verified by
|
|
# eye. A label carrying no digit cannot be mistaken for a dose; a long cell is
|
|
# content rather than a label.
|
|
_DIGIT = re.compile(r"\d")
|
|
_DOSE_VALUE = re.compile(
|
|
r"\d+(?:[.,]\d+)?\s*(?:mg|g|ml|microgam|mcg|µg|%|đơn vị|iu)\b",
|
|
re.IGNORECASE,
|
|
)
|
|
_SUBGROUP_LABEL = re.compile(
|
|
r"\b(?:trẻ|người lớn|người cao tuổi|sơ sinh|thiếu tháng|bệnh nhân|"
|
|
r"suy thận|suy gan|clcr|cân nặng)\b",
|
|
re.IGNORECASE,
|
|
)
|
|
HEADER_CELL_MAX_CHARS = 40
|
|
|
|
|
|
def _is_label_row(cells: Sequence[str]) -> bool:
|
|
kept = [c for c in cells if c and c.strip()]
|
|
if not kept:
|
|
return False
|
|
return all(
|
|
not _DIGIT.search(cell) and len(cell.strip()) <= HEADER_CELL_MAX_CHARS
|
|
for cell in kept
|
|
)
|
|
|
|
|
|
def _split_trailing_label(atom: str) -> List[str]:
|
|
"""Separate a label suffix from clinical text that precedes it.
|
|
|
|
The sentence splitter deliberately ends atoms at ``:``. In list-like dose
|
|
prose that can yield ``"7,5 mg ... .\nBước 5:"`` as one atom. Treating the
|
|
whole atom as the new label makes the dose at its beginning lose ``Bước 4``
|
|
when copied across a seam. Split only at an explicit newline, sentence, or
|
|
semicolon boundary and preserve every character exactly.
|
|
"""
|
|
stripped = atom.rstrip()
|
|
if not stripped.endswith(":"):
|
|
return [atom]
|
|
|
|
body = stripped[:-1]
|
|
cuts = []
|
|
newline = body.rfind("\n")
|
|
if newline >= 0:
|
|
cuts.append(newline + 1)
|
|
for marker in (". ", "; "):
|
|
position = body.rfind(marker)
|
|
if position >= 0:
|
|
cuts.append(position + len(marker))
|
|
if not cuts:
|
|
return [atom]
|
|
|
|
cut = max(cuts)
|
|
if not atom[cut:].strip():
|
|
return [atom]
|
|
return [atom[:cut], atom[cut:]]
|
|
|
|
|
|
def _label_has_embedded_content(atom: str) -> bool:
|
|
"""Whether a colon-ending atom contains content before its final label."""
|
|
prefix = atom.rstrip()[:-1]
|
|
return (
|
|
"\n" in prefix
|
|
or ". " in prefix
|
|
or "; " in prefix
|
|
or bool(_DOSE_VALUE.search(prefix))
|
|
)
|
|
|
|
|
|
def _is_subgroup_label(atom: str) -> bool:
|
|
"""Population/organ-function labels that can sit under a route/indication."""
|
|
return bool(_SUBGROUP_LABEL.search(atom))
|
|
|
|
|
|
def _atoms(text: str, measure: TokenCounter) -> List[str]:
|
|
"""Smallest units the packer may not split.
|
|
|
|
Normally a sentence. A drug-interaction list, though, is one "sentence"
|
|
hundreds of names long: VORICONAZOL's `tương tác thuốc` produced two parts
|
|
of 981 and 888 tokens even after sentence packing. Left that size the
|
|
embedding truncates them, and a truncated interaction list reads as "this
|
|
drug is not listed" — a false negative in exactly the direction that
|
|
matters. Such a run is comma-separated by construction, so a comma is a
|
|
lossless place to break it.
|
|
"""
|
|
atoms: List[str] = []
|
|
sentences = [
|
|
atom
|
|
for sentence in split_sentences(text)
|
|
for atom in _split_trailing_label(sentence)
|
|
]
|
|
for sentence in sentences:
|
|
# Split against the packing target, not the ceiling: an atom sized
|
|
# right up to the ceiling leaves no room for the overlap prepended to
|
|
# it, which is how a 710-token atom became a 981-token chunk.
|
|
if measure(sentence) <= TARGET_TOKENS or "," not in sentence:
|
|
atoms.append(sentence)
|
|
continue
|
|
piece = ""
|
|
for fragment in sentence.split(","):
|
|
candidate = f"{piece},{fragment}" if piece else fragment
|
|
if piece and measure(candidate) > TARGET_TOKENS:
|
|
atoms.append(piece + ",")
|
|
piece = fragment
|
|
else:
|
|
piece = candidate
|
|
if piece:
|
|
atoms.append(piece)
|
|
return atoms
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class _PackedPart:
|
|
atoms: List[str]
|
|
text: str
|
|
source_text: str
|
|
context_labels: List[str]
|
|
|
|
|
|
def _pack_parts(sentences: List[str], measure: TokenCounter) -> List[_PackedPart]:
|
|
"""Greedily pack sentences up to TARGET_TOKENS, overlapping by OVERLAP_TOKENS.
|
|
|
|
A single sentence longer than the target becomes its own part rather than
|
|
being cut mid-sentence — the caller flags it instead of splitting it.
|
|
"""
|
|
IndexedAtom = tuple[int, str]
|
|
PackedAtom = tuple[int, str, bool] # index, text, retrieval-only context
|
|
parts: List[List[PackedAtom]] = []
|
|
current: List[PackedAtom] = []
|
|
current_tokens = 0
|
|
|
|
def is_label(atom: str) -> bool:
|
|
return atom.rstrip().endswith(":")
|
|
|
|
def contains(items: Sequence[PackedAtom], atom: IndexedAtom) -> bool:
|
|
return any((index, text) == atom for index, text, _ in items)
|
|
|
|
# Record the active label for every atom before packing. Looking only inside
|
|
# the current window is insufficient: a population may span several parts,
|
|
# so its label can have fallen outside both the 650-token buffer and the
|
|
# 65-token overlap by the time another dose reaches a seam.
|
|
contexts: List[IndexedAtom | None] = []
|
|
scope_contexts: List[IndexedAtom | None] = []
|
|
active_label: IndexedAtom | None = None
|
|
scope_label: IndexedAtom | None = None
|
|
for index, atom in enumerate(sentences):
|
|
# Context at the *start* of an atom comes from the preceding label. An
|
|
# atom may contain a dose and only end with the next label (Bisoprolol:
|
|
# "7,5 mg ... Bước 5:"); assigning that atom to its own final label
|
|
# makes the dose at its beginning context-free.
|
|
contexts.append(active_label)
|
|
next_scope = scope_label
|
|
if is_label(atom):
|
|
if _is_subgroup_label(atom):
|
|
if active_label is not None and not _is_subgroup_label(active_label[1]):
|
|
next_scope = active_label
|
|
else:
|
|
next_scope = None
|
|
# A pure subgroup label needs its parent scope when it itself lands at
|
|
# the start of a continuation, not only when the following dose arrives.
|
|
scope_contexts.append(next_scope)
|
|
if is_label(atom):
|
|
scope_label = next_scope
|
|
active_label = (index, atom)
|
|
|
|
def context_chain(index: int) -> List[IndexedAtom]:
|
|
"""Labels needed to make atom ``index`` independently interpretable.
|
|
|
|
Some malformed sentence atoms contain a dose and only end with the next
|
|
label. In that case the atom itself is context for what follows, but it
|
|
still needs its own preceding label. Walk only until a plain label;
|
|
this keeps the chain clinically complete without dragging the entire
|
|
section into every overlap.
|
|
"""
|
|
chain: List[IndexedAtom] = []
|
|
|
|
def add(atom: IndexedAtom | None) -> None:
|
|
if atom is not None and atom not in chain:
|
|
chain.append(atom)
|
|
|
|
add(scope_contexts[index])
|
|
|
|
target = sentences[index]
|
|
# A plain label introduces new sibling context, so it needs only its
|
|
# retained parent scope. A compound colon-ending atom still contains
|
|
# clinical material before that new label and therefore needs the
|
|
# preceding active label as well.
|
|
if is_label(target) and not _label_has_embedded_content(target):
|
|
return chain
|
|
|
|
nested: List[IndexedAtom] = []
|
|
context = contexts[index]
|
|
while context is not None and context not in nested:
|
|
nested.append(context)
|
|
add(scope_contexts[context[0]])
|
|
if not _label_has_embedded_content(context[1]):
|
|
break
|
|
context = contexts[context[0]]
|
|
for atom in reversed(nested):
|
|
add(atom)
|
|
return chain
|
|
|
|
for index, sentence in enumerate(sentences):
|
|
tokens = measure(sentence)
|
|
if current and current_tokens + tokens > TARGET_TOKENS:
|
|
# Never end a part on a label. `split_sentences` treats ':' as a
|
|
# boundary, so "Người lớn: 500 mg mỗi 8 giờ." splits after the
|
|
# colon; flushing there leaves a chunk ending "Người lớn:" with
|
|
# the dose in the next one. Measured before this rule: 38 chunks,
|
|
# AMOXICILIN's ending on a Lyme-disease indication followed by a
|
|
# bare "Người lớn:". Outlier item 17 counted population markers on
|
|
# 1,121 of ~1,400 monograph pages, so this is the common shape,
|
|
# and a dose separated from the population it applies to is a
|
|
# patient-safety defect rather than a cosmetic one.
|
|
carried: List[PackedAtom] = []
|
|
while current and is_label(current[-1][1]):
|
|
carried.insert(0, current.pop())
|
|
if not current:
|
|
# A label is safety context, not a useful standalone retrieval
|
|
# unit. Keep it with the following atom even when that makes a
|
|
# synthetic pathological atom oversized; the caller will flag
|
|
# the oversize instead of publishing an unqualified dose.
|
|
current = carried
|
|
current_tokens = sum(measure(atom) for _, atom, _ in current)
|
|
else:
|
|
parts.append(current)
|
|
|
|
overlap: List[PackedAtom] = []
|
|
acc = 0
|
|
for prev_index, prev, is_context in reversed(current):
|
|
# Stop *before* exceeding the budget. A governing label
|
|
# longer than the budget is retained by itself because
|
|
# clinical context wins over the overlap target.
|
|
size = measure(prev)
|
|
if acc + size > OVERLAP_TOKENS:
|
|
break
|
|
overlap.insert(0, (prev_index, prev, is_context))
|
|
acc += size
|
|
|
|
# If dose/detail atoms are copied, their own governing label
|
|
# must precede them even when a newer trailing label is being
|
|
# carried for the incoming sentence.
|
|
if overlap:
|
|
prefix = [
|
|
(context_index, context_text, True)
|
|
for context_index, context_text in context_chain(overlap[0][0])
|
|
if not contains(overlap, (context_index, context_text))
|
|
]
|
|
overlap = prefix + overlap
|
|
current = overlap + carried
|
|
current_tokens = sum(measure(atom) for _, atom, _ in current)
|
|
|
|
# The incoming atom may belong to a label that was flushed several
|
|
# chunks ago. Repeat that label before the new material. Do not add
|
|
# it twice when the incoming atom is itself the label or it was
|
|
# already carried/overlapped.
|
|
incoming_scope = scope_contexts[index]
|
|
for incoming_context in context_chain(index):
|
|
if not contains(current, incoming_context):
|
|
packed_context = (*incoming_context, True)
|
|
if incoming_context == incoming_scope:
|
|
current.insert(0, packed_context)
|
|
else:
|
|
current.append(packed_context)
|
|
current_tokens += measure(incoming_context[1])
|
|
|
|
current.append((index, sentence, False))
|
|
current_tokens += tokens
|
|
|
|
if current:
|
|
parts.append(current)
|
|
return [
|
|
_PackedPart(
|
|
atoms=[atom for _, atom, _ in part],
|
|
text="".join(atom for _, atom, _ in part).strip(),
|
|
source_text="".join(
|
|
atom for _, atom, is_context in part if not is_context
|
|
).strip(),
|
|
context_labels=[
|
|
atom.strip() for _, atom, is_context in part if is_context
|
|
],
|
|
)
|
|
for part in parts
|
|
]
|
|
|
|
|
|
def _pack(sentences: List[str], measure: TokenCounter) -> List[List[str]]:
|
|
"""Compatibility surface used by focused packer tests."""
|
|
return [part.atoms for part in _pack_parts(sentences, measure)]
|
|
|
|
|
|
def _block_kind(block: TableBlock) -> str:
|
|
return KIND_FORMULA if block.shape == SHAPE_FORMULA_2D else KIND_TABLE
|
|
|
|
|
|
def _attachment(
|
|
block: TableBlock,
|
|
_header_row: List[str],
|
|
printed_page: int | None = None,
|
|
) -> ChunkAttachment:
|
|
return ChunkAttachment(
|
|
block_id=block.table_id,
|
|
kind=_block_kind(block),
|
|
shape=block.shape,
|
|
physical_page=block.physical_page,
|
|
bbox=list(block.bbox),
|
|
printed_page=printed_page,
|
|
quarantined=block.quarantined,
|
|
# Header extraction is not human-verified and continuation tables can
|
|
# begin with a body row. Two ARSENIC TRIOXYD ADR rows were previously
|
|
# shipped as "Cột:" metadata. Keep every cell value out of both the
|
|
# retrieval text and vector payload until a reviewed logical-table
|
|
# artifact can prove which row is a header.
|
|
header_row=[],
|
|
)
|
|
|
|
|
|
def _blocks_by_section(monograph: Monograph) -> Dict[str, List[TableBlock]]:
|
|
grouped: Dict[str, List[TableBlock]] = {}
|
|
for block in monograph.tables:
|
|
if block.section_key:
|
|
grouped.setdefault(block.section_key, []).append(block)
|
|
return grouped
|
|
|
|
|
|
def describe_block(monograph: Monograph, section: SectionSpan,
|
|
attachment: ChunkAttachment,
|
|
printed_page: int | None = None) -> str:
|
|
"""Retrieval text for a block, built only from metadata.
|
|
|
|
No cell value ever appears here. A header row is a row of labels;
|
|
linearising it cannot invent a numeric relationship, which is exactly what
|
|
linearising a body row does.
|
|
"""
|
|
noun = "công thức" if attachment.kind == KIND_FORMULA else "bảng"
|
|
page_label = (f"trang {printed_page}" if printed_page is not None
|
|
else "chưa xác định trang in")
|
|
text = (f"{monograph.drug_name} — {section.display_name} — {noun}, "
|
|
f"{page_label}.")
|
|
text += (" Nội dung chỉ tra cứu được trên ảnh trang gốc, "
|
|
"không trích dẫn được dưới dạng văn bản.")
|
|
return text
|
|
|
|
|
|
def _supporting_pages(section: SectionSpan, source_text: str) -> List[int]:
|
|
"""Physical pages whose prose parts intersect one contiguous chunk span."""
|
|
section_text = section.text.strip()
|
|
if not source_text or section_text.count(source_text) != 1:
|
|
raise ValueError(
|
|
f"cannot map {section.key!r} chunk source text uniquely to its section"
|
|
)
|
|
chunk_start = section_text.index(source_text)
|
|
chunk_end = chunk_start + len(source_text)
|
|
|
|
pages: List[int] = []
|
|
cursor = 0
|
|
for part in section.parts:
|
|
if part.kind != "prose" or part.quarantined or not part.text:
|
|
continue
|
|
part_start = section_text.find(part.text, cursor)
|
|
if part_start < 0:
|
|
raise ValueError(
|
|
f"cannot map {section.key!r} part on physical page "
|
|
f"{part.physical_page} back to section text"
|
|
)
|
|
part_end = part_start + len(part.text)
|
|
cursor = part_end
|
|
if part_start < chunk_end and part_end > chunk_start:
|
|
pages.append(part.physical_page)
|
|
|
|
if not pages:
|
|
raise ValueError(f"no page provenance supports section {section.key!r} chunk")
|
|
return sorted(set(pages))
|
|
|
|
|
|
def _page_ranges(
|
|
physical_pages: Sequence[int],
|
|
printed_page_map: Mapping[int, Optional[int]] | None,
|
|
label: str,
|
|
) -> tuple[List[int], List[int]]:
|
|
physical_range = [min(physical_pages), max(physical_pages)]
|
|
if printed_page_map is None:
|
|
return physical_range, []
|
|
printed_pages = [printed_page_map.get(page) for page in physical_pages]
|
|
if any(page is None for page in printed_pages):
|
|
missing = [physical for physical, printed in zip(physical_pages, printed_pages, strict=True)
|
|
if printed is None]
|
|
raise ValueError(
|
|
f"cannot cite {label}: printed folio missing for physical pages {missing}"
|
|
)
|
|
verified = [page for page in printed_pages if page is not None]
|
|
return physical_range, [min(verified), max(verified)]
|
|
|
|
|
|
def chunk_section(monograph: Monograph, section: SectionSpan,
|
|
blocks: Sequence[TableBlock] = (),
|
|
header_rows: Dict[str, List[str]] | None = None,
|
|
printed_page_map: Mapping[int, Optional[int]] | None = None,
|
|
measure: TokenCounter = count_tokens) -> List[Chunk]:
|
|
header_rows = header_rows or {}
|
|
attachments = []
|
|
for block in blocks:
|
|
_, attachment_printed_range = _page_ranges(
|
|
[block.physical_page], printed_page_map,
|
|
f"{monograph.drug_id}/{block.table_id}",
|
|
)
|
|
attachments.append(_attachment(
|
|
block,
|
|
header_rows.get(block.table_id, []),
|
|
attachment_printed_range[0] if attachment_printed_range else None,
|
|
))
|
|
quarantined = any(a.quarantined for a in attachments)
|
|
|
|
def build(body: str, source_body: str, context_labels: List[str],
|
|
part_index: int, part_count: int) -> Chunk:
|
|
tokens = measure(body)
|
|
source_page_range, printed_page_range = _page_ranges(
|
|
_supporting_pages(section, source_body), printed_page_map,
|
|
f"{monograph.drug_id}/{section.key}/{part_index}",
|
|
)
|
|
return Chunk(
|
|
chunk_id=f"{monograph.drug_id}__{section.key}__{part_index}",
|
|
drug_id=monograph.drug_id,
|
|
drug_name=monograph.drug_name,
|
|
section_key=section.key,
|
|
section_display_name=section.display_name,
|
|
text=body,
|
|
source_text=source_body,
|
|
context_labels=list(context_labels),
|
|
heading_physical_page=section.heading.physical_page,
|
|
source_page_range=source_page_range,
|
|
printed_page_range=list(printed_page_range),
|
|
atc_codes=list(monograph.atc_codes),
|
|
part_index=part_index,
|
|
part_count=part_count,
|
|
est_tokens=tokens,
|
|
oversized=tokens > CEILING_TOKENS,
|
|
chunk_kind=CHUNK_KIND_PROSE,
|
|
attachments=list(attachments),
|
|
has_quarantined_content=quarantined,
|
|
)
|
|
|
|
text = section.text.strip()
|
|
prose: List[Chunk] = []
|
|
if text:
|
|
if measure(text) <= CEILING_TOKENS:
|
|
prose = [build(text, text, [], 0, 1)]
|
|
else:
|
|
parts = [part for part in _pack_parts(_atoms(text, measure), measure)
|
|
if part.text]
|
|
prose = [
|
|
build(part.text, part.source_text, part.context_labels,
|
|
index, len(parts))
|
|
for index, part in enumerate(parts)
|
|
]
|
|
|
|
descriptors = []
|
|
for attachment in attachments:
|
|
source_page_range, attachment_printed_range = _page_ranges(
|
|
[attachment.physical_page], printed_page_map,
|
|
f"{monograph.drug_id}/{attachment.block_id}",
|
|
)
|
|
printed_page = (attachment_printed_range[0]
|
|
if attachment_printed_range else None)
|
|
body = describe_block(monograph, section, attachment, printed_page)
|
|
descriptors.append(Chunk(
|
|
chunk_id=f"{monograph.drug_id}__{section.key}__block__{attachment.block_id}",
|
|
drug_id=monograph.drug_id,
|
|
drug_name=monograph.drug_name,
|
|
section_key=section.key,
|
|
section_display_name=section.display_name,
|
|
text=body,
|
|
source_text=body,
|
|
heading_physical_page=section.heading.physical_page,
|
|
source_page_range=source_page_range,
|
|
printed_page_range=attachment_printed_range,
|
|
atc_codes=list(monograph.atc_codes),
|
|
est_tokens=measure(body),
|
|
chunk_kind=CHUNK_KIND_BLOCK_DESCRIPTOR,
|
|
attachments=[attachment],
|
|
has_quarantined_content=attachment.quarantined,
|
|
))
|
|
return prose + descriptors
|
|
|
|
|
|
def chunk_monograph(monograph: Monograph,
|
|
header_rows: Dict[str, List[str]] | None = None,
|
|
printed_page_map: Mapping[int, Optional[int]] | None = None,
|
|
measure: TokenCounter = count_tokens) -> List[Chunk]:
|
|
grouped = _blocks_by_section(monograph)
|
|
chunks: List[Chunk] = []
|
|
for section in monograph.sections.values():
|
|
chunks.extend(chunk_section(monograph, section,
|
|
grouped.get(section.key, ()), header_rows,
|
|
printed_page_map, measure))
|
|
return chunks
|
|
|
|
|
|
def chunk_all(monographs: Iterable[Monograph],
|
|
header_rows: Dict[str, List[str]] | None = None,
|
|
*,
|
|
printed_page_map: Mapping[int, Optional[int]] | None,
|
|
measure: TokenCounter = count_tokens) -> Iterator[Chunk]:
|
|
if printed_page_map is None:
|
|
raise ValueError(
|
|
"chunk_all requires a verified printed_page_map; refusing to emit "
|
|
"an embedding corpus without printed-page provenance"
|
|
)
|
|
for monograph in monographs:
|
|
yield from chunk_monograph(monograph, header_rows, printed_page_map, measure)
|