428 lines
17 KiB
Python
428 lines
17 KiB
Python
import json
|
|
from pathlib import Path
|
|
|
|
from ingestion.chunk import (
|
|
CHUNK_KIND_BLOCK_DESCRIPTOR,
|
|
CHUNK_KIND_PROSE,
|
|
SCHEMA_VERSION,
|
|
chunk_monograph,
|
|
chunk_all,
|
|
chunk_section,
|
|
write_chunks_jsonl,
|
|
)
|
|
from ingestion.chunk.chunker import _is_label_row, describe_block
|
|
from ingestion.segment.models import (
|
|
PART_PROSE,
|
|
Heading,
|
|
Monograph,
|
|
SectionPart,
|
|
SectionSpan,
|
|
TableBlock,
|
|
)
|
|
from ingestion.tables import SHAPE_FORMULA_2D, SHAPE_MULTI_HEADER, SHAPE_SIMPLE
|
|
|
|
|
|
def _section(key, display, text, page=202):
|
|
return SectionSpan(
|
|
key=key, display_name=display,
|
|
heading=Heading(text=display, physical_page=page, y0=100.0,
|
|
is_monograph_title=False, section_key=key),
|
|
text=text,
|
|
parts=([
|
|
SectionPart(
|
|
kind=PART_PROSE,
|
|
text=text,
|
|
physical_page=page,
|
|
bbox=[50.0, 120.0, 550.0, 700.0],
|
|
source_span_ids=[f"p{page}_s0"],
|
|
)
|
|
] if text else []),
|
|
)
|
|
|
|
|
|
def _monograph(sections, tables=()):
|
|
return Monograph(
|
|
drug_id="ampicilin_va_sulbactam",
|
|
drug_name="AMPICILIN VÀ SULBACTAM",
|
|
source_page_range=[200, 203],
|
|
sections={s.key: s for s in sections},
|
|
atc_codes=["J01CR01"],
|
|
tables=list(tables),
|
|
)
|
|
|
|
|
|
def _block(block_id="p202_t0", shape=SHAPE_SIMPLE, section_key="lieu_luong_va_cach_dung"):
|
|
return TableBlock(
|
|
table_id=block_id, shape=shape, physical_page=202,
|
|
bbox=[299.0, 189.6, 552.4, 300.5], section_key=section_key,
|
|
text="Độ thanh thải creatinin Nửa đời Liều 1,5 - 3,0 g",
|
|
quarantined=True,
|
|
)
|
|
|
|
|
|
def test_a_section_whose_table_was_lifted_says_so():
|
|
"""The defect this exists to prevent is silent, not visible.
|
|
|
|
Without the reference, this chunk is grammatical, complete-looking prose
|
|
with the renal-dosing table absent and nothing marking the absence.
|
|
"""
|
|
section = _section("lieu_luong_va_cach_dung", "Liều lượng và cách dùng",
|
|
"Liều thường dùng cho người lớn là 1,5 - 3 g mỗi 6 giờ.")
|
|
monograph = _monograph([section], [_block()])
|
|
chunks = chunk_monograph(monograph)
|
|
|
|
prose = [c for c in chunks if c.chunk_kind == CHUNK_KIND_PROSE]
|
|
assert len(prose) == 1
|
|
assert prose[0].has_quarantined_content is True
|
|
assert [a.block_id for a in prose[0].attachments] == ["p202_t0"]
|
|
assert prose[0].attachments[0].physical_page == 202
|
|
assert prose[0].attachments[0].bbox
|
|
|
|
|
|
def test_a_lifted_block_gets_its_own_retrievable_descriptor():
|
|
section = _section("lieu_luong_va_cach_dung", "Liều lượng và cách dùng", "Prose.")
|
|
monograph = _monograph([section], [_block()])
|
|
descriptors = [c for c in chunk_monograph(
|
|
monograph, printed_page_map={200: 201, 202: 203, 203: 204})
|
|
if c.chunk_kind == CHUNK_KIND_BLOCK_DESCRIPTOR]
|
|
assert len(descriptors) == 1
|
|
assert "AMPICILIN VÀ SULBACTAM" in descriptors[0].text
|
|
assert "Liều lượng và cách dùng" in descriptors[0].text
|
|
# printed page, which is what a reader holding the book looks for
|
|
assert "trang 203" in descriptors[0].text
|
|
assert descriptors[0].source_page_range == [202, 202]
|
|
assert descriptors[0].printed_page_range == [203, 203]
|
|
assert descriptors[0].attachments[0].printed_page == 203
|
|
|
|
|
|
def test_no_cell_value_ever_reaches_the_descriptor_text():
|
|
section = _section("lieu_luong_va_cach_dung", "Liều lượng và cách dùng", "Prose.")
|
|
block = _block()
|
|
monograph = _monograph([section], [block])
|
|
descriptors = [c for c in chunk_monograph(monograph, {"p202_t0": []})
|
|
if c.chunk_kind == CHUNK_KIND_BLOCK_DESCRIPTOR]
|
|
assert "1,5 - 3,0 g" not in descriptors[0].text
|
|
|
|
|
|
def test_a_header_row_carrying_a_number_is_refused():
|
|
"""AMIODARON, physical page 183 — a real case, caught by a gate.
|
|
|
|
pdfplumber reported the first row as
|
|
"Thời gian liệu pháp tĩnh mạch Liều 720 mg/ngày (0,5 mg/phút)", i.e. a
|
|
dose inside what it called a header, from an extraction never verified by
|
|
eye. Measured: 42 of 124 simple-table headers (34%) contain a digit.
|
|
"""
|
|
assert _is_label_row(["Các Statin", "Khởi đầu", "Liều duy trì"]) is True
|
|
assert _is_label_row(["Liều 720 mg/ngày (0,5 mg/phút)"]) is False
|
|
assert _is_label_row(["x" * 45]) is False
|
|
assert _is_label_row([]) is False
|
|
|
|
section = _section("lieu_luong_va_cach_dung", "Liều lượng và cách dùng", "Prose.")
|
|
monograph = _monograph([section], [_block()])
|
|
chunks = chunk_monograph(
|
|
monograph, {"p202_t0": ["Liều 720 mg/ngày (0,5 mg/phút)"]})
|
|
descriptor = next(c for c in chunks
|
|
if c.chunk_kind == CHUNK_KIND_BLOCK_DESCRIPTOR)
|
|
assert "720" not in descriptor.text
|
|
assert descriptor.attachments[0].header_row == []
|
|
|
|
|
|
def test_unverified_header_rows_are_embargoed_for_every_table_shape():
|
|
section = _section("lieu_luong_va_cach_dung", "Liều lượng và cách dùng", "Prose.")
|
|
header = {"p202_t0": ["Ngoại tâm thu thất", "Thường gặp", "Không rõ tần suất"]}
|
|
for shape in (SHAPE_SIMPLE, SHAPE_MULTI_HEADER):
|
|
monograph = _monograph([section], [_block(shape=shape)])
|
|
descriptor = next(c for c in chunk_monograph(monograph, header)
|
|
if c.chunk_kind == CHUNK_KIND_BLOCK_DESCRIPTOR)
|
|
assert descriptor.attachments[0].header_row == []
|
|
assert "Cột:" not in descriptor.text
|
|
assert "Ngoại tâm thu thất" not in descriptor.text
|
|
assert "Không rõ tần suất" not in descriptor.text
|
|
|
|
|
|
def test_a_formula_block_is_described_as_a_formula():
|
|
section = _section("than_trong", "Thận trọng", "Prose.")
|
|
block = _block(block_id="p1042_f0", shape=SHAPE_FORMULA_2D,
|
|
section_key="than_trong")
|
|
monograph = _monograph([section], [block])
|
|
descriptor = next(c for c in chunk_monograph(monograph)
|
|
if c.chunk_kind == CHUNK_KIND_BLOCK_DESCRIPTOR)
|
|
assert "công thức" in descriptor.text
|
|
assert "bảng" not in descriptor.text
|
|
|
|
|
|
def test_attachments_do_not_change_the_prose_text():
|
|
"""The condition under which this feature was accepted at all."""
|
|
section = _section("lieu_luong_va_cach_dung", "Liều lượng và cách dùng",
|
|
"Liều thường dùng cho người lớn là 1,5 - 3 g mỗi 6 giờ.")
|
|
with_block = chunk_section(_monograph([section], [_block()]), section,
|
|
[_block()])
|
|
without = chunk_section(_monograph([section]), section)
|
|
prose_with = [c for c in with_block if c.chunk_kind == CHUNK_KIND_PROSE]
|
|
assert [c.text for c in prose_with] == [c.text for c in without]
|
|
assert [c.chunk_id for c in prose_with] == [c.chunk_id for c in without]
|
|
|
|
|
|
def test_a_section_with_no_text_but_a_block_still_yields_the_descriptor():
|
|
section = _section("lieu_luong_va_cach_dung", "Liều lượng và cách dùng", "")
|
|
chunks = chunk_monograph(_monograph([section], [_block()]))
|
|
assert [c.chunk_kind for c in chunks] == [CHUNK_KIND_BLOCK_DESCRIPTOR]
|
|
|
|
|
|
def test_written_chunks_declare_their_schema_version(tmp_path: Path):
|
|
section = _section("chi_dinh", "Chỉ định", "Nhiễm khuẩn.")
|
|
chunks = chunk_monograph(_monograph([section]))
|
|
out = tmp_path / "chunks.jsonl"
|
|
assert write_chunks_jsonl(chunks, out) == 1
|
|
record = json.loads(out.read_text(encoding="utf-8").splitlines()[0])
|
|
assert record["schema_version"] == SCHEMA_VERSION
|
|
assert record["chunk_kind"] == CHUNK_KIND_PROSE
|
|
assert record["has_quarantined_content"] is False
|
|
|
|
|
|
def test_describe_block_names_the_page_even_with_no_header():
|
|
section = _section("lieu_luong_va_cach_dung", "Liều lượng và cách dùng", "Prose.")
|
|
monograph = _monograph([section], [_block()])
|
|
chunks = chunk_monograph(monograph)
|
|
attachment = next(c for c in chunks
|
|
if c.chunk_kind == CHUNK_KIND_BLOCK_DESCRIPTOR).attachments[0]
|
|
text = describe_block(monograph, section, attachment, printed_page=203)
|
|
assert "trang 203" in text
|
|
assert "không trích dẫn được dưới dạng văn bản" in text
|
|
|
|
|
|
def test_chunk_carries_only_verified_printed_page_range():
|
|
section = _section("chi_dinh", "Chỉ định", "Nhiễm khuẩn.")
|
|
monograph = _monograph([section])
|
|
chunk = chunk_monograph(
|
|
monograph,
|
|
printed_page_map={202: 203},
|
|
)[0]
|
|
assert chunk.source_page_range == [202, 202]
|
|
assert chunk.printed_page_range == [203, 203]
|
|
|
|
|
|
def test_chunk_refuses_an_unmapped_printed_folio():
|
|
section = _section("chi_dinh", "Chỉ định", "Nhiễm khuẩn.")
|
|
monograph = _monograph([section])
|
|
try:
|
|
chunk_monograph(monograph, printed_page_map={202: None})
|
|
except ValueError as exc:
|
|
assert "printed folio missing" in str(exc)
|
|
else:
|
|
raise AssertionError("missing printed folio must fail closed")
|
|
|
|
|
|
def test_a_chunk_spanning_two_section_parts_cites_only_those_pages():
|
|
text = "Nội dung trang một.\nNội dung trang hai."
|
|
section = SectionSpan(
|
|
key="chi_dinh",
|
|
display_name="Chỉ định",
|
|
heading=Heading(
|
|
text="Chỉ định", physical_page=201, y0=100.0,
|
|
is_monograph_title=False, section_key="chi_dinh",
|
|
),
|
|
text=text,
|
|
parts=[
|
|
SectionPart(PART_PROSE, "Nội dung trang một.", 201,
|
|
[50.0, 100.0, 550.0, 200.0], ["p201_s0"]),
|
|
SectionPart(PART_PROSE, "Nội dung trang hai.", 202,
|
|
[50.0, 100.0, 550.0, 200.0], ["p202_s0"]),
|
|
],
|
|
)
|
|
chunk = chunk_monograph(
|
|
_monograph([section]), printed_page_map={201: 202, 202: 203}
|
|
)[0]
|
|
|
|
assert chunk.source_page_range == [201, 202]
|
|
assert chunk.printed_page_range == [202, 203]
|
|
|
|
|
|
def test_whole_corpus_chunking_refuses_to_run_without_a_printed_page_map():
|
|
section = _section("chi_dinh", "Chỉ định", "Nhiễm khuẩn.")
|
|
try:
|
|
list(chunk_all([_monograph([section])], printed_page_map=None))
|
|
except ValueError as exc:
|
|
assert "requires a verified printed_page_map" in str(exc)
|
|
else:
|
|
raise AssertionError("whole-corpus chunking must fail closed without folios")
|
|
|
|
|
|
def test_the_char_ratio_estimate_is_never_used_as_a_token_count():
|
|
"""ADR 0004 sized chunks with len(text)//4 and reported 0 over the ceiling.
|
|
|
|
Counted with the real tokenizer, 1,884 of 12,838 chunks (14.7%) were over
|
|
it, the largest at 1,645 tokens — twice the ceiling. Vietnamese diacritics
|
|
cost multiple byte-pair tokens each; measured ratio real/estimate is 1.95
|
|
at the median and 6.0 at worst.
|
|
"""
|
|
from ingestion.chunk.tokens import count_tokens, estimate_tokens
|
|
|
|
vietnamese = "Liều thường dùng cho người lớn là 1,5 - 3 g mỗi 6 giờ."
|
|
assert count_tokens(vietnamese) > len(vietnamese) // 4
|
|
# the fallback errs small, so it can never certify an oversized chunk as safe
|
|
assert estimate_tokens(vietnamese) > len(vietnamese) // 4
|
|
|
|
|
|
def test_a_long_comma_list_is_split_at_commas_not_left_oversized():
|
|
"""VORICONAZOL's interaction list is one 'sentence' hundreds of names long.
|
|
|
|
Truncated by an embedding model it reads as "this drug is not listed" — a
|
|
false negative in the direction that matters. A comma is a lossless break.
|
|
"""
|
|
from ingestion.chunk.chunker import CEILING_TOKENS, _atoms
|
|
|
|
drugs = ", ".join(f"thuốc {n}" for n in range(400))
|
|
atoms = _atoms(drugs + ".", lambda t: len(t) // 2)
|
|
assert len(atoms) > 1
|
|
assert all(len(a) // 2 <= CEILING_TOKENS for a in atoms)
|
|
assert "".join(atoms).replace(",", "") == (drugs + ".").replace(",", "")
|
|
|
|
|
|
def test_the_overlap_never_exceeds_its_budget():
|
|
"""A 251-token atom produced a 273-token overlap against a 65-token
|
|
setting, because the loop added whole atoms until the total passed it.
|
|
That was most of how a 981-token chunk came about."""
|
|
from ingestion.chunk.chunker import OVERLAP_TOKENS, _pack
|
|
|
|
measure = lambda t: len(t) # noqa: E731 - one-line stub for the test
|
|
atoms = ["a" * 300, "b" * 300, "c" * 300]
|
|
parts = _pack(atoms, measure)
|
|
assert len(parts) > 1
|
|
for part in parts[1:]:
|
|
carried = part[:-1]
|
|
assert sum(measure(a) for a in carried) <= OVERLAP_TOKENS
|
|
|
|
|
|
def test_a_continuation_repeats_the_label_governing_its_dose():
|
|
"""A budget-only overlap used to strand population labels.
|
|
|
|
The two 30-token dose atoms fit the 65-token overlap, while the preceding
|
|
label did not. The continuation was therefore independently retrievable
|
|
as a bare dose even though its source context was population-specific.
|
|
"""
|
|
from ingestion.chunk.chunker import _pack
|
|
|
|
measure = len
|
|
label = "Trẻ đẻ thiếu tháng:"
|
|
atoms = [
|
|
"p" * 570,
|
|
label,
|
|
"Uống liều 2 mg/kg q12h. " + "a" * 5,
|
|
"Nếu không uống được: " + "b" * 8,
|
|
"Theo dõi đáp ứng và điều chỉnh liều. " + "c" * 70,
|
|
]
|
|
|
|
parts = _pack(atoms, measure)
|
|
|
|
assert len(parts) == 2
|
|
assert label in parts[1]
|
|
assert parts[1].index(label) < next(
|
|
index for index, atom in enumerate(parts[1]) if "liều" in atom
|
|
)
|
|
|
|
|
|
def test_a_single_long_label_is_never_emitted_without_its_dose():
|
|
"""When the current buffer held only one long label, the old loop could
|
|
not carry it and emitted a label-only retrievable chunk."""
|
|
from ingestion.chunk.chunker import _pack
|
|
|
|
label = "Trẻ sơ sinh có tình trạng lâm sàng cần hiệu chỉnh đặc biệt " * 2 + ":"
|
|
dose = "Dùng liều khởi đầu " + "x" * 640
|
|
parts = _pack([label, dose], len)
|
|
|
|
assert all(part != [label] for part in parts)
|
|
assert any(label in part and dose in part for part in parts)
|
|
|
|
|
|
def test_a_new_trailing_label_does_not_orphan_the_previous_population_dose():
|
|
"""Real shape: a neonatal dose is followed by ``Suy thận:`` at the seam.
|
|
|
|
Carrying only the new trailing label is insufficient: any dose atoms copied
|
|
into the overlap must retain the older population label that governs them.
|
|
"""
|
|
from ingestion.chunk.chunker import _pack
|
|
|
|
population = "Trẻ đẻ thiếu tháng và trẻ sơ sinh dưới 8 ngày tuổi:"
|
|
renal = "Suy thận:"
|
|
first_dose = "100 mg/kg/ngày, chia hai lần. "
|
|
dose_limit = "Liều tối đa 10 mg/kg/ngày. "
|
|
atoms = [
|
|
"p" * 520,
|
|
population,
|
|
first_dose,
|
|
dose_limit,
|
|
renal,
|
|
"Điều chỉnh theo độ thanh thải creatinin. " + "x" * 80,
|
|
]
|
|
|
|
parts = _pack(atoms, len)
|
|
|
|
assert len(parts) == 2
|
|
assert parts[1][0] == population
|
|
assert renal in parts[1]
|
|
copied_doses = [atom for atom in parts[1] if atom in (first_dose, dose_limit)]
|
|
if copied_doses:
|
|
assert parts[1].index(population) < min(parts[1].index(atom) for atom in copied_doses)
|
|
|
|
|
|
def test_an_atom_ending_in_the_next_label_keeps_the_previous_dose_context():
|
|
"""Bisoprolol has atoms shaped ``dose for step 4 ... Step 5:``.
|
|
|
|
Ending in a colon does not make the dose at the beginning of that same atom
|
|
belong to the new label.
|
|
"""
|
|
from ingestion.chunk.chunker import _pack, _split_trailing_label
|
|
|
|
previous = "Bước 4:"
|
|
compound = "7,5 mg/lần/ngày trong 4 tuần; chuyển bước 5.\nBước 5:"
|
|
next_dose = "10 mg/lần/ngày để duy trì. "
|
|
split_compound = _split_trailing_label(compound)
|
|
assert "".join(split_compound) == compound
|
|
assert split_compound == [
|
|
"7,5 mg/lần/ngày trong 4 tuần; chuyển bước 5.\n",
|
|
"Bước 5:",
|
|
]
|
|
atoms = [
|
|
"p" * 540,
|
|
previous,
|
|
*split_compound,
|
|
next_dose,
|
|
"Theo dõi dung nạp và điều chỉnh. " + "x" * 80,
|
|
]
|
|
|
|
parts = _pack(atoms, len)
|
|
|
|
assert len(parts) == 2
|
|
dose_atom = split_compound[0]
|
|
if dose_atom in parts[1]:
|
|
assert previous in parts[1]
|
|
assert parts[1].index(previous) < parts[1].index(dose_atom)
|
|
assert split_compound[1] in parts[1]
|
|
assert parts[1].index(split_compound[1]) < parts[1].index(next_dose)
|
|
|
|
|
|
def test_a_population_continuation_retains_its_parent_route():
|
|
"""PARACETAMOL: age-band labels are children of ``Đường trực tràng:``."""
|
|
from ingestion.chunk.chunker import _pack
|
|
|
|
route = "Đường trực tràng:"
|
|
population = "Trẻ em 1 - 3 tháng tuổi:"
|
|
dose = "30 mg/kg một liều duy nhất. "
|
|
next_population = "Trẻ em 3 tháng - 6 tuổi:"
|
|
atoms = [
|
|
"p" * 540,
|
|
route,
|
|
population,
|
|
dose,
|
|
next_population,
|
|
"30 - 40 mg/kg một liều duy nhất. " + "x" * 80,
|
|
]
|
|
|
|
parts = _pack(atoms, len)
|
|
|
|
assert len(parts) == 2
|
|
assert route in parts[1]
|
|
assert next_population in parts[1]
|
|
assert parts[1].index(route) < parts[1].index(next_population)
|