169 lines
7.3 KiB
Python
169 lines
7.3 KiB
Python
import json
|
|
from pathlib import Path
|
|
|
|
from ingestion.chunk import (
|
|
CHUNK_KIND_BLOCK_DESCRIPTOR,
|
|
CHUNK_KIND_PROSE,
|
|
SCHEMA_VERSION,
|
|
chunk_monograph,
|
|
chunk_section,
|
|
write_chunks_jsonl,
|
|
)
|
|
from ingestion.chunk.chunker import _is_label_row, describe_block
|
|
from ingestion.segment.models import Heading, Monograph, 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,
|
|
)
|
|
|
|
|
|
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)
|
|
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
|
|
|
|
|
|
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_only_a_simple_table_contributes_a_header():
|
|
section = _section("lieu_luong_va_cach_dung", "Liều lượng và cách dùng", "Prose.")
|
|
header = {"p202_t0": ["Nhóm", "Liều"]}
|
|
for shape, expected in ((SHAPE_SIMPLE, ["Nhóm", "Liều"]),
|
|
(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 == expected
|
|
|
|
|
|
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)
|
|
assert "trang 203" in text
|
|
assert "không trích dẫn được dưới dạng văn bản" in text
|