Wire the guarded conversational RAG answer layer end-to-end
This commit is contained in:
@@ -0,0 +1,17 @@
|
||||
{
|
||||
"source_id": "dtqgvn_2_2018",
|
||||
"title": "Dược thư Quốc gia Việt Nam - lần xuất bản thứ hai",
|
||||
"edition": 2,
|
||||
"publication_year": 2018,
|
||||
"pdf_path": "data/raw/duoc-thu-quoc-gia-viet-nam-2018.pdf",
|
||||
"sha256": "2aa81c846a5e760f82658c46816ab63174204a7d95b3e2755b6565288e53d0d1",
|
||||
"superseded_by": {
|
||||
"edition": 3,
|
||||
"publication_year": 2022,
|
||||
"decision": "3445/QĐ-BYT",
|
||||
"decision_date": "2022-12-23"
|
||||
},
|
||||
"production_use_rights_documented": false,
|
||||
"clinical_production_eligible": false,
|
||||
"note": "Suitable for parser development and historical comparison only; not the sole source for clinical production."
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,4 +1,5 @@
|
||||
from .chunker import chunk_all, chunk_monograph, chunk_section, estimate_tokens
|
||||
from .chunker import chunk_all, chunk_monograph, chunk_section
|
||||
from .tokens import count_tokens, estimate_tokens, tokenizer_available
|
||||
from .io import read_monographs_jsonl, write_chunks_jsonl
|
||||
from .models import (
|
||||
CHUNK_KIND_BLOCK_DESCRIPTOR,
|
||||
@@ -17,7 +18,9 @@ __all__ = [
|
||||
"chunk_all",
|
||||
"chunk_monograph",
|
||||
"chunk_section",
|
||||
"count_tokens",
|
||||
"estimate_tokens",
|
||||
"tokenizer_available",
|
||||
"read_monographs_jsonl",
|
||||
"write_chunks_jsonl",
|
||||
]
|
||||
|
||||
@@ -7,10 +7,11 @@ sub-chunked, with a sentence-boundary-aware sliding window.
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from typing import Dict, Iterable, Iterator, List, Sequence
|
||||
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, SHAPE_SIMPLE
|
||||
from ..tables.classify import SHAPE_FORMULA_2D
|
||||
from .models import (
|
||||
CHUNK_KIND_BLOCK_DESCRIPTOR,
|
||||
CHUNK_KIND_PROSE,
|
||||
@@ -18,16 +19,12 @@ from .models import (
|
||||
ChunkAttachment,
|
||||
)
|
||||
from .sentences import split_sentences
|
||||
from .tokens import TokenCounter, count_tokens
|
||||
|
||||
CEILING_TOKENS = 800
|
||||
TARGET_TOKENS = 650
|
||||
OVERLAP_TOKENS = 65
|
||||
|
||||
# Physical -> printed page. Empirically constant across every tested
|
||||
# milestone page (extract/page_map.py, ADR 0003); the descriptor quotes the
|
||||
# printed number because that is what a reader holding the book looks for.
|
||||
PRINTED_PAGE_OFFSET = 1
|
||||
|
||||
KIND_TABLE = "table"
|
||||
KIND_FORMULA = "formula"
|
||||
|
||||
@@ -39,6 +36,15 @@ KIND_FORMULA = "formula"
|
||||
# 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
|
||||
|
||||
|
||||
@@ -52,61 +58,293 @@ def _is_label_row(cells: Sequence[str]) -> bool:
|
||||
)
|
||||
|
||||
|
||||
def estimate_tokens(text: str) -> int:
|
||||
"""ADR 0004's chars/4 estimate — an estimate, not a tokenizer count."""
|
||||
return len(text) // 4
|
||||
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 _pack(sentences: List[str]) -> List[List[str]]:
|
||||
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.
|
||||
"""
|
||||
parts: List[List[str]] = []
|
||||
current: List[str] = []
|
||||
IndexedAtom = tuple[int, str]
|
||||
PackedAtom = tuple[int, str, bool] # index, text, retrieval-only context
|
||||
parts: List[List[PackedAtom]] = []
|
||||
current: List[PackedAtom] = []
|
||||
current_tokens = 0
|
||||
|
||||
for sentence in sentences:
|
||||
tokens = estimate_tokens(sentence)
|
||||
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:
|
||||
parts.append(current)
|
||||
overlap: List[str] = []
|
||||
acc = 0
|
||||
for prev in reversed(current):
|
||||
overlap.insert(0, prev)
|
||||
acc += estimate_tokens(prev)
|
||||
if acc >= OVERLAP_TOKENS:
|
||||
break
|
||||
current = list(overlap)
|
||||
current_tokens = sum(estimate_tokens(s) for s in current)
|
||||
current.append(sentence)
|
||||
# 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 parts
|
||||
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]) -> ChunkAttachment:
|
||||
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,
|
||||
# Only a simple table's first row can be a row of plain labels, and
|
||||
# only when it actually reads like one. A multi-level or merged header
|
||||
# is the shape whose extraction is least trustworthy, so it
|
||||
# contributes nothing rather than something wrong.
|
||||
header_row=(list(header_row)
|
||||
if block.shape == SHAPE_SIMPLE and _is_label_row(header_row)
|
||||
else []),
|
||||
# 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=[],
|
||||
)
|
||||
|
||||
|
||||
@@ -119,7 +357,8 @@ def _blocks_by_section(monograph: Monograph) -> Dict[str, List[TableBlock]]:
|
||||
|
||||
|
||||
def describe_block(monograph: Monograph, section: SectionSpan,
|
||||
attachment: ChunkAttachment) -> str:
|
||||
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;
|
||||
@@ -127,28 +366,91 @@ def describe_block(monograph: Monograph, section: SectionSpan,
|
||||
linearising a body row does.
|
||||
"""
|
||||
noun = "công thức" if attachment.kind == KIND_FORMULA else "bảng"
|
||||
printed = attachment.physical_page + PRINTED_PAGE_OFFSET
|
||||
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"trang {printed}.")
|
||||
if attachment.header_row:
|
||||
columns = " | ".join(c.replace("\n", " ").strip()
|
||||
for c in attachment.header_row if c and c.strip())
|
||||
if columns:
|
||||
text += f" Cột: {columns}."
|
||||
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) -> List[Chunk]:
|
||||
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 = [_attachment(b, header_rows.get(b.table_id, [])) for b in blocks]
|
||||
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, part_index: int, part_count: int) -> Chunk:
|
||||
tokens = estimate_tokens(body)
|
||||
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,
|
||||
@@ -156,8 +458,11 @@ def chunk_section(monograph: Monograph, section: SectionSpan,
|
||||
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=list(monograph.source_page_range),
|
||||
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,
|
||||
@@ -171,16 +476,26 @@ def chunk_section(monograph: Monograph, section: SectionSpan,
|
||||
text = section.text.strip()
|
||||
prose: List[Chunk] = []
|
||||
if text:
|
||||
if estimate_tokens(text) <= CEILING_TOKENS:
|
||||
prose = [build(text, 0, 1)]
|
||||
if measure(text) <= CEILING_TOKENS:
|
||||
prose = [build(text, text, [], 0, 1)]
|
||||
else:
|
||||
parts = _pack(split_sentences(text))
|
||||
bodies = [b for b in ("".join(p).strip() for p in parts) if b]
|
||||
prose = [build(b, i, len(bodies)) for i, b in enumerate(bodies)]
|
||||
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:
|
||||
body = describe_block(monograph, section, attachment)
|
||||
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,
|
||||
@@ -188,10 +503,12 @@ def chunk_section(monograph: Monograph, section: SectionSpan,
|
||||
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=list(monograph.source_page_range),
|
||||
source_page_range=source_page_range,
|
||||
printed_page_range=attachment_printed_range,
|
||||
atc_codes=list(monograph.atc_codes),
|
||||
est_tokens=estimate_tokens(body),
|
||||
est_tokens=measure(body),
|
||||
chunk_kind=CHUNK_KIND_BLOCK_DESCRIPTOR,
|
||||
attachments=[attachment],
|
||||
has_quarantined_content=attachment.quarantined,
|
||||
@@ -200,16 +517,27 @@ def chunk_section(monograph: Monograph, section: SectionSpan,
|
||||
|
||||
|
||||
def chunk_monograph(monograph: Monograph,
|
||||
header_rows: Dict[str, List[str]] | None = None) -> List[Chunk]:
|
||||
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))
|
||||
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) -> Iterator[Chunk]:
|
||||
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)
|
||||
yield from chunk_monograph(monograph, header_rows, printed_page_map, measure)
|
||||
|
||||
@@ -6,7 +6,13 @@ from dataclasses import asdict
|
||||
from pathlib import Path
|
||||
from typing import Iterable, Iterator
|
||||
|
||||
from ..segment.models import Heading, Monograph, SectionSpan, TableBlock
|
||||
from ..segment.models import (
|
||||
Heading,
|
||||
Monograph,
|
||||
SectionPart,
|
||||
SectionSpan,
|
||||
TableBlock,
|
||||
)
|
||||
from .models import SCHEMA_VERSION, Chunk
|
||||
|
||||
|
||||
@@ -31,6 +37,23 @@ def read_monographs_jsonl(path: Path) -> Iterator[Monograph]:
|
||||
section_key=h.get("section_key"),
|
||||
),
|
||||
text=s["text"],
|
||||
# Carried, not dropped: CLAUDE.md's provenance rule is
|
||||
# explicit that a stage boundary must not shed fields.
|
||||
parts=[
|
||||
SectionPart(
|
||||
kind=p["kind"],
|
||||
text=p["text"],
|
||||
physical_page=p["physical_page"],
|
||||
bbox=p["bbox"],
|
||||
source_span_ids=p.get("source_span_ids", []),
|
||||
table_id=p.get("table_id"),
|
||||
table_part_id=p.get("table_part_id"),
|
||||
continuation_group=p.get("continuation_group"),
|
||||
shape=p.get("shape"),
|
||||
quarantined=p.get("quarantined", False),
|
||||
)
|
||||
for p in s.get("parts", [])
|
||||
],
|
||||
)
|
||||
yield Monograph(
|
||||
drug_id=raw["drug_id"],
|
||||
|
||||
@@ -18,7 +18,7 @@ from __future__ import annotations
|
||||
from dataclasses import dataclass, field
|
||||
from typing import List
|
||||
|
||||
SCHEMA_VERSION = 2
|
||||
SCHEMA_VERSION = 4
|
||||
|
||||
CHUNK_KIND_PROSE = "prose"
|
||||
CHUNK_KIND_BLOCK_DESCRIPTOR = "block_descriptor"
|
||||
@@ -37,6 +37,8 @@ class ChunkAttachment:
|
||||
shape: str
|
||||
physical_page: int
|
||||
bbox: List[float]
|
||||
printed_page: int | None = None
|
||||
source_crop: str | None = None
|
||||
quarantined: bool = True
|
||||
# First row of a `simple_table`, used to make the block findable. Comes
|
||||
# from pdfplumber and has NOT been verified by eye — the 180 real tables'
|
||||
@@ -54,6 +56,12 @@ class Chunk:
|
||||
text: str
|
||||
heading_physical_page: int
|
||||
source_page_range: List[int]
|
||||
# Exact contiguous source material. `text` may prepend retrieval-only
|
||||
# context labels at a seam; provenance/reassembly must never mistake those
|
||||
# repetitions for a literal source span.
|
||||
source_text: str = ""
|
||||
context_labels: List[str] = field(default_factory=list)
|
||||
printed_page_range: List[int] = field(default_factory=list)
|
||||
atc_codes: List[str] = field(default_factory=list)
|
||||
part_index: int = 0
|
||||
part_count: int = 1
|
||||
|
||||
@@ -0,0 +1,65 @@
|
||||
"""Token counting for the chunk stage.
|
||||
|
||||
ADR 0004 sized chunks with `len(text) // 4`, describing it honestly as an
|
||||
estimate. Measured against the real tokenizer on this corpus, that estimate is
|
||||
wrong by about a factor of two for Vietnamese: real/estimate is **1.95 at the
|
||||
median, 2.50 at p95, 6.0 at worst**, because accented Vietnamese characters
|
||||
cost multiple byte-pair tokens each where English prose costs about four
|
||||
characters per token.
|
||||
|
||||
The consequence was not academic. Under the estimate the pipeline reported
|
||||
**0 chunks over the 800-token ceiling**; counted properly, **1,884 of 12,838
|
||||
(14.7%)** were over it, the largest at 1,645 tokens — twice the ceiling. A
|
||||
reassuring number that was simply false.
|
||||
|
||||
The counter is injectable so the chunking logic stays testable without the
|
||||
tokenizer installed, and so a different embedding model's tokenizer can be
|
||||
substituted without touching chunk shaping.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Callable
|
||||
|
||||
# What OpenAI's text-embedding-3-* models use. Named here rather than inside
|
||||
# the function so swapping models is a one-line, visible change.
|
||||
ENCODING_NAME = "cl100k_base"
|
||||
|
||||
# Only for the no-tokenizer fallback. Derived from the measurement above
|
||||
# (median 1.95 real tokens per chars/4 unit), i.e. ~2 characters per token —
|
||||
# still an estimate, but one that errs on the side of smaller chunks instead
|
||||
# of larger ones.
|
||||
FALLBACK_CHARS_PER_TOKEN = 2
|
||||
|
||||
TokenCounter = Callable[[str], int]
|
||||
|
||||
_encoder = None
|
||||
|
||||
|
||||
def _load_encoder():
|
||||
global _encoder
|
||||
if _encoder is None:
|
||||
import tiktoken
|
||||
|
||||
_encoder = tiktoken.get_encoding(ENCODING_NAME)
|
||||
return _encoder
|
||||
|
||||
|
||||
def count_tokens(text: str) -> int:
|
||||
"""Real token count, falling back to a conservative estimate."""
|
||||
try:
|
||||
return len(_load_encoder().encode(text))
|
||||
except Exception:
|
||||
return estimate_tokens(text)
|
||||
|
||||
|
||||
def estimate_tokens(text: str) -> int:
|
||||
"""Character-ratio fallback. An estimate — never report it as a count."""
|
||||
return len(text) // FALLBACK_CHARS_PER_TOKEN
|
||||
|
||||
|
||||
def tokenizer_available() -> bool:
|
||||
try:
|
||||
_load_encoder()
|
||||
return True
|
||||
except Exception:
|
||||
return False
|
||||
@@ -16,6 +16,7 @@ from pathlib import Path
|
||||
import fitz
|
||||
|
||||
from .extract import (
|
||||
build_page_map,
|
||||
extract_spans,
|
||||
load_transcribed_runs,
|
||||
merge_outlined_runs,
|
||||
@@ -26,6 +27,7 @@ from .extract import (
|
||||
)
|
||||
from .chunk import (
|
||||
CHUNK_KIND_PROSE,
|
||||
tokenizer_available,
|
||||
chunk_all,
|
||||
read_monographs_jsonl,
|
||||
write_chunks_jsonl,
|
||||
@@ -151,9 +153,12 @@ def _cmd_validate(args: argparse.Namespace) -> int:
|
||||
return 1
|
||||
|
||||
doc = fitz.open(pdf_path)
|
||||
spans = list(extract_spans(doc))
|
||||
# Same stream and same regions as `run`. Measuring recall against a
|
||||
# pipeline that is not the one producing the output is how `coverage`
|
||||
# ended up describing a different build earlier today.
|
||||
spans = _extracted_and_repaired_spans(doc)
|
||||
try:
|
||||
monographs = list(assemble(spans))
|
||||
monographs = list(assemble(spans, table_index=_region_index(args.tables)))
|
||||
except DuplicateDrugIdError as e:
|
||||
print(f"error: {e}", file=sys.stderr)
|
||||
return 1
|
||||
@@ -332,8 +337,17 @@ def _cmd_chunk(args: argparse.Namespace) -> int:
|
||||
header_rows = {r.table_id: r.first_row
|
||||
for r in read_regions_json(regions_path)}
|
||||
|
||||
pdf_path = Path(args.pdf)
|
||||
if not pdf_path.exists():
|
||||
print(f"error: no source PDF at {pdf_path}", file=sys.stderr)
|
||||
return 1
|
||||
with fitz.open(pdf_path) as document:
|
||||
printed_page_map = build_page_map(document)
|
||||
|
||||
monographs = list(read_monographs_jsonl(monographs_path))
|
||||
chunks = list(chunk_all(monographs, header_rows))
|
||||
chunks = list(chunk_all(
|
||||
monographs, header_rows, printed_page_map=printed_page_map,
|
||||
))
|
||||
|
||||
kinds = Counter(c.chunk_kind for c in chunks)
|
||||
with_attachments = sum(1 for c in chunks
|
||||
@@ -347,7 +361,7 @@ def _cmd_chunk(args: argparse.Namespace) -> int:
|
||||
print(f" {kind:<22}{n:>7}")
|
||||
print(f"prose chunks carrying a lifted block: {with_attachments}")
|
||||
print(f"oversized (over the {800}-token ceiling): {oversized}")
|
||||
print(f"estimated tokens (chars/4, an estimate): {tokens:,}")
|
||||
print(f"tokens ({'cl100k_base' if tokenizer_available() else 'ESTIMATED — tokenizer missing'}): {tokens:,}")
|
||||
|
||||
out_path = Path(args.out)
|
||||
written = write_chunks_jsonl(chunks, out_path)
|
||||
@@ -444,6 +458,7 @@ def build_parser() -> argparse.ArgumentParser:
|
||||
|
||||
p_validate = sub.add_parser("validate", help="Whole-book recall/precision vs. back-of-book index")
|
||||
p_validate.add_argument("--pdf", required=True)
|
||||
p_validate.add_argument("--tables", default="data/processed/table_regions.json")
|
||||
p_validate.set_defaults(func=_cmd_validate)
|
||||
|
||||
p_tables = sub.add_parser(
|
||||
@@ -491,6 +506,8 @@ def build_parser() -> argparse.ArgumentParser:
|
||||
"chunk", help="Build retrieval chunks (ADR 0004/0005/0006)")
|
||||
p_chunk.add_argument("--monographs", default="data/processed/monographs.jsonl")
|
||||
p_chunk.add_argument("--tables", default="data/processed/table_regions.json")
|
||||
p_chunk.add_argument(
|
||||
"--pdf", default="data/raw/duoc-thu-quoc-gia-viet-nam-2018.pdf")
|
||||
p_chunk.add_argument("--out", default="data/processed/chunks.jsonl")
|
||||
p_chunk.set_defaults(func=_cmd_chunk)
|
||||
|
||||
|
||||
@@ -0,0 +1,54 @@
|
||||
"""Embedding stage: text in, vectors out, with the model behind an interface.
|
||||
|
||||
Import from here rather than from a provider module. Nothing outside this
|
||||
package should name boto3, `sentence_transformers`, or a model id — that is
|
||||
what lets `load/` and the retrieval side stay testable without a live account,
|
||||
and what makes swapping the benchmark winner a one-line change.
|
||||
"""
|
||||
from .bedrock_runtime import DEFAULT_REGION, BedrockInvoker, Boto3BedrockInvoker
|
||||
from .cache import (
|
||||
CacheStats,
|
||||
CachingEmbeddingProvider,
|
||||
EmbeddingCache,
|
||||
cache_key,
|
||||
)
|
||||
from .ports import (
|
||||
INPUT_DOCUMENT,
|
||||
INPUT_KINDS,
|
||||
INPUT_QUERY,
|
||||
EmbeddingBatch,
|
||||
EmbeddingProvider,
|
||||
EmbeddingVector,
|
||||
text_digest,
|
||||
)
|
||||
from .registry import (
|
||||
BGE_M3,
|
||||
CLOUD_PROVIDERS,
|
||||
COHERE_V4,
|
||||
TITAN_V2,
|
||||
build_provider,
|
||||
provider_names,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"BGE_M3",
|
||||
"CLOUD_PROVIDERS",
|
||||
"COHERE_V4",
|
||||
"DEFAULT_REGION",
|
||||
"INPUT_DOCUMENT",
|
||||
"INPUT_KINDS",
|
||||
"INPUT_QUERY",
|
||||
"TITAN_V2",
|
||||
"BedrockInvoker",
|
||||
"Boto3BedrockInvoker",
|
||||
"CacheStats",
|
||||
"CachingEmbeddingProvider",
|
||||
"EmbeddingBatch",
|
||||
"EmbeddingCache",
|
||||
"EmbeddingProvider",
|
||||
"EmbeddingVector",
|
||||
"build_provider",
|
||||
"cache_key",
|
||||
"provider_names",
|
||||
"text_digest",
|
||||
]
|
||||
|
||||
@@ -0,0 +1,152 @@
|
||||
"""Cohere Embed v4 (`cohere.embed-v4:0`).
|
||||
|
||||
Request/response shape taken from the AWS Bedrock user guide page "Cohere
|
||||
Embed v4" (read 2026-08-03):
|
||||
|
||||
request {"input_type": "search_document|search_query|classification|
|
||||
clustering",
|
||||
"texts": [str], # max 96 per call
|
||||
"embedding_types": ["float"|"int8"|"uint8"|"binary"|"ubinary"],
|
||||
"output_dimension": 256|512|1024|1536,
|
||||
"truncate": "NONE|LEFT|RIGHT"}
|
||||
|
||||
The response has two documented shapes and this adapter accepts both. Asking
|
||||
for one or more `embedding_types` returns
|
||||
`{"response_type": "embeddings_by_type", "embeddings": {"float": [[...]]}}`;
|
||||
omitting the field returns
|
||||
`{"response_type": "embeddings_floats", "embeddings": [[...]]}`.
|
||||
|
||||
Three defaults are chosen rather than inherited:
|
||||
|
||||
- `output_dimension` is set explicitly. The documented default is 1536, and a
|
||||
collection built at one width cannot absorb vectors of another.
|
||||
- `input_type` is derived from the caller's input kind. This is the model
|
||||
whose asymmetry the `ports` contract exists for: corpus records go in as
|
||||
`search_document`, queries as `search_query`.
|
||||
- `truncate` is `NONE`, which makes an over-length input an error instead of a
|
||||
silently shortened one. A dosing section that lost its tail and embedded
|
||||
anyway is exactly the failure this project's rules are written against.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, List, Sequence
|
||||
|
||||
from .bedrock_runtime import BedrockInvoker
|
||||
from .ports import (
|
||||
INPUT_DOCUMENT,
|
||||
INPUT_QUERY,
|
||||
EmbeddingProvider,
|
||||
EmbeddingVector,
|
||||
text_digest,
|
||||
)
|
||||
|
||||
MODEL_ID = "cohere.embed-v4:0"
|
||||
PROVIDER_NAME = "cohere-v4"
|
||||
|
||||
SUPPORTED_DIMENSIONS = (256, 512, 1024, 1536)
|
||||
|
||||
# The documented per-request ceiling for `texts`.
|
||||
MAX_TEXTS_PER_REQUEST = 96
|
||||
|
||||
COHERE_INPUT_TYPES = {
|
||||
INPUT_DOCUMENT: "search_document",
|
||||
INPUT_QUERY: "search_query",
|
||||
}
|
||||
|
||||
# Cohere's docs do not state whether float vectors are unit-length, so this
|
||||
# stays unset rather than being asserted either way.
|
||||
NORMALIZED_UNKNOWN = None
|
||||
|
||||
|
||||
class CohereEmbedV4(EmbeddingProvider):
|
||||
def __init__(
|
||||
self,
|
||||
invoker: BedrockInvoker,
|
||||
dimensions: int = 1024,
|
||||
truncate: str = "NONE",
|
||||
batch_size: int = MAX_TEXTS_PER_REQUEST,
|
||||
):
|
||||
if dimensions not in SUPPORTED_DIMENSIONS:
|
||||
raise ValueError(
|
||||
f"{MODEL_ID} supports {SUPPORTED_DIMENSIONS}, got {dimensions}"
|
||||
)
|
||||
if not 1 <= batch_size <= MAX_TEXTS_PER_REQUEST:
|
||||
raise ValueError(
|
||||
f"batch_size must be 1..{MAX_TEXTS_PER_REQUEST}, got {batch_size}"
|
||||
)
|
||||
self._invoker = invoker
|
||||
self._dimensions = dimensions
|
||||
self._truncate = truncate
|
||||
self._batch_size = batch_size
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return PROVIDER_NAME
|
||||
|
||||
@property
|
||||
def model_id(self) -> str:
|
||||
return MODEL_ID
|
||||
|
||||
@property
|
||||
def dimensions(self) -> int:
|
||||
return self._dimensions
|
||||
|
||||
@property
|
||||
def max_batch_size(self) -> int:
|
||||
return self._batch_size
|
||||
|
||||
def _embed_batch(
|
||||
self, texts: Sequence[str], input_kind: str
|
||||
) -> List[EmbeddingVector]:
|
||||
body = self._invoker.invoke_json(
|
||||
MODEL_ID,
|
||||
{
|
||||
"texts": list(texts),
|
||||
"input_type": COHERE_INPUT_TYPES[input_kind],
|
||||
"embedding_types": ["float"],
|
||||
"output_dimension": self._dimensions,
|
||||
"truncate": self._truncate,
|
||||
},
|
||||
# The AWS code example for this model sends `*/*`.
|
||||
accept="*/*",
|
||||
)
|
||||
rows = _float_rows(body)
|
||||
if len(rows) != len(texts):
|
||||
raise ValueError(
|
||||
f"{MODEL_ID} returned {len(rows)} vectors for {len(texts)} texts"
|
||||
)
|
||||
|
||||
vectors: List[EmbeddingVector] = []
|
||||
for text, values in zip(texts, rows, strict=True):
|
||||
self._check_dimensions(values)
|
||||
vectors.append(
|
||||
EmbeddingVector(
|
||||
values=list(values),
|
||||
text_sha256=text_digest(text),
|
||||
provider=PROVIDER_NAME,
|
||||
model_id=MODEL_ID,
|
||||
dimensions=self._dimensions,
|
||||
input_kind=input_kind,
|
||||
normalized=NORMALIZED_UNKNOWN,
|
||||
)
|
||||
)
|
||||
return vectors
|
||||
|
||||
|
||||
def _float_rows(body: Any) -> List[List[float]]:
|
||||
"""Pull the float vectors out of either documented response shape."""
|
||||
embeddings = body.get("embeddings")
|
||||
if embeddings is None:
|
||||
raise ValueError(
|
||||
f"{MODEL_ID} response has no 'embeddings' field; "
|
||||
f"keys were {sorted(body)}"
|
||||
)
|
||||
if isinstance(embeddings, dict):
|
||||
rows = embeddings.get("float")
|
||||
if rows is None:
|
||||
raise ValueError(
|
||||
f"{MODEL_ID} returned no float embeddings; "
|
||||
f"types present: {sorted(embeddings)}"
|
||||
)
|
||||
return rows
|
||||
return embeddings
|
||||
@@ -0,0 +1,75 @@
|
||||
"""The only module in this package that knows boto3 exists.
|
||||
|
||||
Keeping the SDK behind `BedrockInvoker` is what makes the two Bedrock
|
||||
adapters testable with no AWS account, no credentials and no spend: a test
|
||||
passes a stub that returns a canned response body, and the adapter's request
|
||||
shaping and response parsing — the parts that can actually be wrong — are
|
||||
exercised in full.
|
||||
|
||||
boto3 is imported inside the method rather than at module scope so that
|
||||
`ingestion.embed` imports cleanly in an environment that never talks to AWS
|
||||
(the local bge-m3 control, or the mocked tests).
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from typing import Any, Mapping, Optional, Protocol
|
||||
|
||||
DEFAULT_REGION = "us-east-1"
|
||||
|
||||
BEDROCK_RUNTIME_SERVICE = "bedrock-runtime"
|
||||
|
||||
|
||||
class BedrockInvoker(Protocol):
|
||||
def invoke_json(
|
||||
self,
|
||||
model_id: str,
|
||||
payload: Mapping[str, Any],
|
||||
accept: str = "application/json",
|
||||
) -> dict:
|
||||
"""POST `payload` as JSON to a Bedrock model, return the parsed body."""
|
||||
|
||||
|
||||
class Boto3BedrockInvoker:
|
||||
"""`InvokeModel` over boto3, with the JSON/stream plumbing hidden."""
|
||||
|
||||
def __init__(self, region: str = DEFAULT_REGION, client: Optional[Any] = None):
|
||||
self._region = region
|
||||
self._client = client
|
||||
|
||||
@property
|
||||
def region(self) -> str:
|
||||
return self._region
|
||||
|
||||
def _runtime(self) -> Any:
|
||||
if self._client is None:
|
||||
import boto3
|
||||
from botocore.config import Config
|
||||
|
||||
# Without an explicit read timeout a single stalled response hangs
|
||||
# the whole corpus run: observed 2026-08-04, one request held an
|
||||
# open socket for over five minutes while boto3's default waited.
|
||||
self._client = boto3.client(
|
||||
BEDROCK_RUNTIME_SERVICE,
|
||||
region_name=self._region,
|
||||
config=Config(
|
||||
connect_timeout=10,
|
||||
read_timeout=60,
|
||||
retries={"max_attempts": 3, "mode": "standard"},
|
||||
),
|
||||
)
|
||||
return self._client
|
||||
|
||||
def invoke_json(
|
||||
self,
|
||||
model_id: str,
|
||||
payload: Mapping[str, Any],
|
||||
accept: str = "application/json",
|
||||
) -> dict:
|
||||
response = self._runtime().invoke_model(
|
||||
modelId=model_id,
|
||||
body=json.dumps(payload),
|
||||
accept=accept,
|
||||
contentType="application/json",
|
||||
)
|
||||
return json.loads(response["body"].read())
|
||||
@@ -0,0 +1,96 @@
|
||||
"""Amazon Titan Text Embeddings V2 (`amazon.titan-embed-text-v2:0`).
|
||||
|
||||
Request/response shape taken from the AWS Bedrock user guide page
|
||||
"Amazon Titan Embeddings G1 - Text", V2 tabs (read 2026-08-03):
|
||||
|
||||
request {"inputText": str, "dimensions": int, "normalize": bool,
|
||||
"embeddingTypes": list}
|
||||
response {"embedding": [float], "inputTextTokenCount": int,
|
||||
"embeddingsByType": {...}}
|
||||
|
||||
`embeddingTypes` is left unset so the response keeps the plain float
|
||||
`embedding` field — the documentation notes that field disappears when
|
||||
`embeddingTypes` contains only `binary`.
|
||||
|
||||
Titan draws no distinction between a corpus record and a query, so both input
|
||||
kinds produce a byte-identical request. The kind is still recorded on each
|
||||
vector, because "this model ignores it" is a fact worth being able to read
|
||||
back off the data rather than infer.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import List, Sequence
|
||||
|
||||
from .bedrock_runtime import BedrockInvoker
|
||||
from .ports import EmbeddingProvider, EmbeddingVector, text_digest
|
||||
|
||||
MODEL_ID = "amazon.titan-embed-text-v2:0"
|
||||
PROVIDER_NAME = "titan-v2"
|
||||
|
||||
# Per the V2 request documentation.
|
||||
SUPPORTED_DIMENSIONS = (256, 512, 1024)
|
||||
|
||||
|
||||
class TitanTextEmbeddingsV2(EmbeddingProvider):
|
||||
def __init__(
|
||||
self,
|
||||
invoker: BedrockInvoker,
|
||||
dimensions: int = 1024,
|
||||
normalize: bool = True,
|
||||
):
|
||||
if dimensions not in SUPPORTED_DIMENSIONS:
|
||||
raise ValueError(
|
||||
f"{MODEL_ID} supports {SUPPORTED_DIMENSIONS}, got {dimensions}"
|
||||
)
|
||||
self._invoker = invoker
|
||||
self._dimensions = dimensions
|
||||
self._normalize = normalize
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return PROVIDER_NAME
|
||||
|
||||
@property
|
||||
def model_id(self) -> str:
|
||||
return MODEL_ID
|
||||
|
||||
@property
|
||||
def dimensions(self) -> int:
|
||||
return self._dimensions
|
||||
|
||||
@property
|
||||
def max_batch_size(self) -> int:
|
||||
# InvokeModel takes a single `inputText`; there is no text array.
|
||||
return 1
|
||||
|
||||
def _embed_batch(
|
||||
self, texts: Sequence[str], input_kind: str
|
||||
) -> List[EmbeddingVector]:
|
||||
text = texts[0]
|
||||
body = self._invoker.invoke_json(
|
||||
MODEL_ID,
|
||||
{
|
||||
"inputText": text,
|
||||
"dimensions": self._dimensions,
|
||||
"normalize": self._normalize,
|
||||
},
|
||||
)
|
||||
values = body.get("embedding")
|
||||
if values is None:
|
||||
raise ValueError(
|
||||
f"{MODEL_ID} response has no 'embedding' field; "
|
||||
f"keys were {sorted(body)}"
|
||||
)
|
||||
self._check_dimensions(values)
|
||||
return [
|
||||
EmbeddingVector(
|
||||
values=list(values),
|
||||
text_sha256=text_digest(text),
|
||||
provider=PROVIDER_NAME,
|
||||
model_id=MODEL_ID,
|
||||
dimensions=self._dimensions,
|
||||
input_kind=input_kind,
|
||||
normalized=self._normalize,
|
||||
input_token_count=body.get("inputTextTokenCount"),
|
||||
)
|
||||
]
|
||||
@@ -0,0 +1,247 @@
|
||||
"""A disk cache so a corpus is never paid for twice.
|
||||
|
||||
**Why the key is content-addressed.** A vector is a pure function of three
|
||||
things: the model, the input kind, and the exact bytes embedded. Nothing else
|
||||
about the record changes the answer. `docs/v1-delivery-plan.md` §4.A proposed
|
||||
keying on `chunk_id` + sha256; measured against the real corpus that would
|
||||
charge twice for identical text — `chunks.jsonl` holds 15,066 records but only
|
||||
14,869 distinct texts, so 197 records (1.31%) are repeats of a text already
|
||||
embedded. The key here is `(model_id, input_kind, text_sha256)`, which collapses
|
||||
those and, more importantly, cannot silently serve a stale vector after a chunk's
|
||||
text is edited: an edit changes the digest, so it is a miss.
|
||||
|
||||
Traceability is not lost by dropping `chunk_id` from the key. Every cached
|
||||
record carries the same sha256 rule (`ports.text_digest`) that produced it, so a
|
||||
vector is matched back to its chunk by re-digesting that chunk's text. Pairing
|
||||
vectors to chunk records is `load/`'s job, not the cache's.
|
||||
|
||||
**Why the index holds offsets, not vectors.** 15,066 vectors of 1,024 floats do
|
||||
not belong in memory all at once; a list of that many Python floats is tens of
|
||||
kilobytes each. Startup scans the file once to map key to byte offset, and a
|
||||
`get` seeks and parses exactly one line.
|
||||
|
||||
The file is append-only. A key already present is never rewritten, so the file
|
||||
is a log that can be inspected, truncated, or resumed after an interrupted run
|
||||
without a repair step.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Dict, Iterable, List, Optional, Sequence, Tuple
|
||||
|
||||
from .ports import (
|
||||
INPUT_KINDS,
|
||||
EmbeddingBatch,
|
||||
EmbeddingProvider,
|
||||
EmbeddingVector,
|
||||
text_digest,
|
||||
)
|
||||
|
||||
CacheKey = Tuple[str, str, str]
|
||||
|
||||
_KEY_FIELDS = ("model_id", "input_kind", "text_sha256")
|
||||
|
||||
|
||||
def cache_key(model_id: str, input_kind: str, text: str) -> CacheKey:
|
||||
return (model_id, input_kind, text_digest(text))
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class CacheStats:
|
||||
hits: int = 0
|
||||
misses: int = 0
|
||||
|
||||
@property
|
||||
def lookups(self) -> int:
|
||||
return self.hits + self.misses
|
||||
|
||||
@property
|
||||
def hit_rate(self) -> float:
|
||||
return self.hits / self.lookups if self.lookups else 0.0
|
||||
|
||||
|
||||
class EmbeddingCache:
|
||||
"""Append-only JSONL of vectors, indexed by byte offset."""
|
||||
|
||||
def __init__(self, path: os.PathLike | str) -> None:
|
||||
self._path = Path(path)
|
||||
self._offsets: Dict[CacheKey, int] = {}
|
||||
self._hits = 0
|
||||
self._misses = 0
|
||||
if self._path.exists():
|
||||
self._build_index()
|
||||
|
||||
@property
|
||||
def path(self) -> Path:
|
||||
return self._path
|
||||
|
||||
@property
|
||||
def stats(self) -> CacheStats:
|
||||
return CacheStats(hits=self._hits, misses=self._misses)
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self._offsets)
|
||||
|
||||
def __contains__(self, key: CacheKey) -> bool:
|
||||
return key in self._offsets
|
||||
|
||||
def _build_index(self) -> None:
|
||||
with self._path.open("rb") as handle:
|
||||
offset = 0
|
||||
for raw in handle:
|
||||
line = raw.decode("utf-8").strip()
|
||||
if line:
|
||||
record = json.loads(line)
|
||||
self._offsets[self._key_of(record)] = offset
|
||||
offset += len(raw)
|
||||
|
||||
@staticmethod
|
||||
def _key_of(record: dict) -> CacheKey:
|
||||
missing = [f for f in _KEY_FIELDS if not record.get(f)]
|
||||
if missing:
|
||||
raise ValueError(f"cache record is missing key fields: {missing}")
|
||||
return (
|
||||
record["model_id"],
|
||||
record["input_kind"],
|
||||
record["text_sha256"],
|
||||
)
|
||||
|
||||
def get(self, key: CacheKey) -> Optional[EmbeddingVector]:
|
||||
offset = self._offsets.get(key)
|
||||
if offset is None:
|
||||
self._misses += 1
|
||||
return None
|
||||
with self._path.open("rb") as handle:
|
||||
handle.seek(offset)
|
||||
record = json.loads(handle.readline().decode("utf-8"))
|
||||
self._hits += 1
|
||||
return _vector_from_record(record)
|
||||
|
||||
def put(self, vector: EmbeddingVector) -> bool:
|
||||
"""Append a vector. Returns False if the key was already stored."""
|
||||
key = (vector.model_id, vector.input_kind, vector.text_sha256)
|
||||
if key in self._offsets:
|
||||
return False
|
||||
self._path.parent.mkdir(parents=True, exist_ok=True)
|
||||
line = json.dumps(_record_from_vector(vector), ensure_ascii=False) + "\n"
|
||||
encoded = line.encode("utf-8")
|
||||
with self._path.open("ab") as handle:
|
||||
offset = handle.tell()
|
||||
handle.write(encoded)
|
||||
self._offsets[key] = offset
|
||||
return True
|
||||
|
||||
def put_many(self, vectors: Iterable[EmbeddingVector]) -> int:
|
||||
return sum(1 for vector in vectors if self.put(vector))
|
||||
|
||||
|
||||
def _record_from_vector(vector: EmbeddingVector) -> dict:
|
||||
return {
|
||||
"model_id": vector.model_id,
|
||||
"input_kind": vector.input_kind,
|
||||
"text_sha256": vector.text_sha256,
|
||||
"provider": vector.provider,
|
||||
"dimensions": vector.dimensions,
|
||||
"normalized": vector.normalized,
|
||||
"input_token_count": vector.input_token_count,
|
||||
"values": vector.values,
|
||||
}
|
||||
|
||||
|
||||
def _vector_from_record(record: dict) -> EmbeddingVector:
|
||||
values: List[float] = record["values"]
|
||||
declared = record["dimensions"]
|
||||
if len(values) != declared:
|
||||
raise ValueError(
|
||||
f"cached vector for {record['text_sha256'][:12]} has {len(values)} "
|
||||
f"values but declares {declared} dimensions"
|
||||
)
|
||||
return EmbeddingVector(
|
||||
values=values,
|
||||
text_sha256=record["text_sha256"],
|
||||
provider=record["provider"],
|
||||
model_id=record["model_id"],
|
||||
dimensions=declared,
|
||||
input_kind=record["input_kind"],
|
||||
normalized=record.get("normalized"),
|
||||
input_token_count=record.get("input_token_count"),
|
||||
)
|
||||
|
||||
|
||||
class CachingEmbeddingProvider(EmbeddingProvider):
|
||||
"""Wraps a provider so only uncached texts reach it.
|
||||
|
||||
A decorator rather than a change to each adapter: the three existing
|
||||
providers stay unaware that a cache exists, and a fourth needs no cache code
|
||||
to benefit. `request_count` counts requests the *inner* provider actually
|
||||
made, which is what makes "a second run costs nothing" a checkable claim
|
||||
rather than an assertion — a fully cached run reports zero.
|
||||
"""
|
||||
|
||||
def __init__(self, inner: EmbeddingProvider, cache: EmbeddingCache) -> None:
|
||||
self._inner = inner
|
||||
self._cache = cache
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return f"cached:{self._inner.name}"
|
||||
|
||||
@property
|
||||
def model_id(self) -> str:
|
||||
return self._inner.model_id
|
||||
|
||||
@property
|
||||
def dimensions(self) -> int:
|
||||
return self._inner.dimensions
|
||||
|
||||
@property
|
||||
def max_batch_size(self) -> int:
|
||||
return self._inner.max_batch_size
|
||||
|
||||
@property
|
||||
def cache(self) -> EmbeddingCache:
|
||||
return self._cache
|
||||
|
||||
def _embed_batch(
|
||||
self, texts: Sequence[str], input_kind: str
|
||||
) -> List[EmbeddingVector]:
|
||||
return list(self._inner.embed(list(texts), input_kind).vectors)
|
||||
|
||||
def embed(self, texts: Sequence[str], input_kind: str) -> EmbeddingBatch:
|
||||
if input_kind not in INPUT_KINDS:
|
||||
raise ValueError(
|
||||
f"input_kind must be one of {INPUT_KINDS}, got {input_kind!r}"
|
||||
)
|
||||
if any(not t.strip() for t in texts):
|
||||
raise ValueError("refusing to embed an empty or whitespace-only text")
|
||||
|
||||
resolved: List[Optional[EmbeddingVector]] = []
|
||||
pending: Dict[str, List[int]] = {}
|
||||
for position, text in enumerate(texts):
|
||||
hit = self._cache.get(cache_key(self.model_id, input_kind, text))
|
||||
resolved.append(hit)
|
||||
if hit is None:
|
||||
pending.setdefault(text, []).append(position)
|
||||
|
||||
requests = 0
|
||||
latency_ms = 0.0
|
||||
if pending:
|
||||
wanted = list(pending)
|
||||
batch = self._inner.embed(wanted, input_kind)
|
||||
requests = batch.request_count
|
||||
latency_ms = batch.latency_ms
|
||||
for text, vector in zip(wanted, batch.vectors, strict=True):
|
||||
self._cache.put(vector)
|
||||
for position in pending[text]:
|
||||
resolved[position] = vector
|
||||
|
||||
if any(vector is None for vector in resolved):
|
||||
raise ValueError("cache resolution left a text without a vector")
|
||||
return EmbeddingBatch(
|
||||
vectors=[vector for vector in resolved if vector is not None],
|
||||
request_count=requests,
|
||||
latency_ms=latency_ms,
|
||||
)
|
||||
@@ -0,0 +1,103 @@
|
||||
"""BAAI/bge-m3 running locally — the zero-API-cost control in the benchmark.
|
||||
|
||||
Its job is to answer "how much is the paid model actually buying us on
|
||||
Vietnamese medical prose?". Without a free baseline in the same harness, a
|
||||
cloud model's recall number has nothing to be better *than*.
|
||||
|
||||
Two properties are taken from the published model card and have **not** been
|
||||
verified on this machine (no local run has happened yet — see the coordination
|
||||
handoff): the dense vector is 1024-dimensional, and bge-m3 needs no
|
||||
instruction prefix on either the corpus or the query side, unlike the earlier
|
||||
English bge models. Both are asserted at runtime rather than trusted: the
|
||||
dimension is checked on every vector by `EmbeddingProvider._check_dimensions`,
|
||||
so a wrong assumption fails on the first call instead of producing a
|
||||
quietly unusable collection.
|
||||
|
||||
`sentence-transformers` is imported lazily and the encoder is injectable, so
|
||||
this module costs nothing to import and can be tested without the ~2 GB of
|
||||
model weights.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Callable, List, Optional, Sequence
|
||||
|
||||
from .ports import EmbeddingProvider, EmbeddingVector, text_digest
|
||||
|
||||
MODEL_ID = "BAAI/bge-m3"
|
||||
PROVIDER_NAME = "bge-m3"
|
||||
|
||||
DENSE_DIMENSIONS = 1024
|
||||
|
||||
Encoder = Callable[[Sequence[str]], Sequence[Sequence[float]]]
|
||||
|
||||
|
||||
class BgeM3Local(EmbeddingProvider):
|
||||
def __init__(
|
||||
self,
|
||||
encoder: Optional[Encoder] = None,
|
||||
batch_size: int = 16,
|
||||
device: Optional[str] = None,
|
||||
):
|
||||
if batch_size < 1:
|
||||
raise ValueError(f"batch_size must be >= 1, got {batch_size}")
|
||||
self._encoder = encoder
|
||||
self._batch_size = batch_size
|
||||
self._device = device
|
||||
# Only the encoder built below is known to normalize. An injected one
|
||||
# is somebody else's function, so its output is recorded as unknown.
|
||||
self._normalized: Optional[bool] = None if encoder is not None else True
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return PROVIDER_NAME
|
||||
|
||||
@property
|
||||
def model_id(self) -> str:
|
||||
return MODEL_ID
|
||||
|
||||
@property
|
||||
def dimensions(self) -> int:
|
||||
return DENSE_DIMENSIONS
|
||||
|
||||
@property
|
||||
def max_batch_size(self) -> int:
|
||||
return self._batch_size
|
||||
|
||||
def _load_encoder(self) -> Encoder:
|
||||
if self._encoder is None:
|
||||
from sentence_transformers import SentenceTransformer
|
||||
|
||||
model = SentenceTransformer(MODEL_ID, device=self._device)
|
||||
|
||||
def encode(texts: Sequence[str]) -> Sequence[Sequence[float]]:
|
||||
return model.encode(
|
||||
list(texts), normalize_embeddings=True
|
||||
).tolist()
|
||||
|
||||
self._encoder = encode
|
||||
return self._encoder
|
||||
|
||||
def _embed_batch(
|
||||
self, texts: Sequence[str], input_kind: str
|
||||
) -> List[EmbeddingVector]:
|
||||
rows = self._load_encoder()(texts)
|
||||
if len(rows) != len(texts):
|
||||
raise ValueError(
|
||||
f"{MODEL_ID} returned {len(rows)} vectors for {len(texts)} texts"
|
||||
)
|
||||
|
||||
vectors: List[EmbeddingVector] = []
|
||||
for text, values in zip(texts, rows, strict=True):
|
||||
self._check_dimensions(values)
|
||||
vectors.append(
|
||||
EmbeddingVector(
|
||||
values=list(values),
|
||||
text_sha256=text_digest(text),
|
||||
provider=PROVIDER_NAME,
|
||||
model_id=MODEL_ID,
|
||||
dimensions=DENSE_DIMENSIONS,
|
||||
input_kind=input_kind,
|
||||
normalized=self._normalized,
|
||||
)
|
||||
)
|
||||
return vectors
|
||||
@@ -0,0 +1,144 @@
|
||||
"""The embedding boundary: what a provider must do, and what it must record.
|
||||
|
||||
Two things drive this design.
|
||||
|
||||
**Asymmetric models make the input kind part of the contract.** Cohere embeds
|
||||
a corpus record and a search query into deliberately different subspaces —
|
||||
the same string sent as `search_document` and as `search_query` does not come
|
||||
back as the same vector. Getting that backwards raises no error; recall just
|
||||
quietly drops. So `input_kind` is a required argument of `embed()`, not an
|
||||
optional keyword a caller can forget, and the value used is recorded on every
|
||||
vector so a mismatch is detectable after the fact.
|
||||
|
||||
**Provenance applies to vectors too.** CLAUDE.md requires an extracted unit to
|
||||
stay traceable to its source; a vector is no different. `model_id`,
|
||||
`dimensions`, `input_kind` and the sha256 of the exact text embedded are what
|
||||
let a collection be checked for the one mistake that is invisible from the
|
||||
outside — vectors from two different models mixed into one Qdrant collection,
|
||||
where every query still returns *something*.
|
||||
|
||||
`normalized` is deliberately three-valued. Titan is asked to normalize and
|
||||
says so; a local encoder is told to; Cohere's Bedrock documentation does not
|
||||
state whether its float vectors are unit-length, so the field stays `None`
|
||||
rather than guessing. An unmeasured claim does not get written down as a fact.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import time
|
||||
from abc import ABC, abstractmethod
|
||||
from dataclasses import dataclass, field
|
||||
from typing import List, Optional, Sequence
|
||||
|
||||
INPUT_DOCUMENT = "document"
|
||||
INPUT_QUERY = "query"
|
||||
|
||||
INPUT_KINDS = (INPUT_DOCUMENT, INPUT_QUERY)
|
||||
|
||||
|
||||
def text_digest(text: str) -> str:
|
||||
"""sha256 of the exact string sent to the provider.
|
||||
|
||||
Shared by every adapter so a cached vector can be matched to its text by
|
||||
the same rule that produced it.
|
||||
"""
|
||||
return hashlib.sha256(text.encode("utf-8")).hexdigest()
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class EmbeddingVector:
|
||||
values: List[float]
|
||||
text_sha256: str
|
||||
provider: str
|
||||
model_id: str
|
||||
dimensions: int
|
||||
input_kind: str
|
||||
# None means the provider does not document it — not "no".
|
||||
normalized: Optional[bool] = None
|
||||
# Only some providers report it (Titan does, Cohere's documented text
|
||||
# response does not).
|
||||
input_token_count: Optional[int] = None
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class EmbeddingBatch:
|
||||
"""Vectors plus what a benchmark needs to compare providers fairly."""
|
||||
|
||||
vectors: List[EmbeddingVector] = field(default_factory=list)
|
||||
request_count: int = 0
|
||||
latency_ms: float = 0.0
|
||||
|
||||
|
||||
class EmbeddingProvider(ABC):
|
||||
"""One embedding model, reachable without the caller knowing its SDK.
|
||||
|
||||
Subclasses implement `_embed_batch` for a single request; `embed` owns
|
||||
input validation, splitting into provider-sized requests, and timing, so
|
||||
that logic exists once rather than per adapter.
|
||||
"""
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def name(self) -> str:
|
||||
"""Short registry key, e.g. `titan-v2`."""
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def model_id(self) -> str:
|
||||
"""Provider-side identifier, e.g. `amazon.titan-embed-text-v2:0`."""
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def dimensions(self) -> int:
|
||||
"""Vector length this instance is configured to produce."""
|
||||
|
||||
@property
|
||||
def max_batch_size(self) -> int:
|
||||
"""Texts accepted per request. Default is the safest possible value."""
|
||||
return 1
|
||||
|
||||
@abstractmethod
|
||||
def _embed_batch(
|
||||
self, texts: Sequence[str], input_kind: str
|
||||
) -> List[EmbeddingVector]:
|
||||
"""Embed at most `max_batch_size` texts in one provider request."""
|
||||
|
||||
def embed(self, texts: Sequence[str], input_kind: str) -> EmbeddingBatch:
|
||||
if input_kind not in INPUT_KINDS:
|
||||
raise ValueError(
|
||||
f"input_kind must be one of {INPUT_KINDS}, got {input_kind!r}"
|
||||
)
|
||||
if any(not t.strip() for t in texts):
|
||||
raise ValueError("refusing to embed an empty or whitespace-only text")
|
||||
|
||||
vectors: List[EmbeddingVector] = []
|
||||
requests = 0
|
||||
started = time.perf_counter()
|
||||
for start in range(0, len(texts), self.max_batch_size):
|
||||
window = texts[start : start + self.max_batch_size]
|
||||
vectors.extend(self._embed_batch(window, input_kind))
|
||||
requests += 1
|
||||
elapsed_ms = (time.perf_counter() - started) * 1000.0
|
||||
|
||||
if len(vectors) != len(texts):
|
||||
raise ValueError(
|
||||
f"{self.name} returned {len(vectors)} vectors for "
|
||||
f"{len(texts)} texts"
|
||||
)
|
||||
return EmbeddingBatch(
|
||||
vectors=vectors, request_count=requests, latency_ms=elapsed_ms
|
||||
)
|
||||
|
||||
def embed_documents(self, texts: Sequence[str]) -> EmbeddingBatch:
|
||||
return self.embed(texts, INPUT_DOCUMENT)
|
||||
|
||||
def embed_queries(self, texts: Sequence[str]) -> EmbeddingBatch:
|
||||
return self.embed(texts, INPUT_QUERY)
|
||||
|
||||
def _check_dimensions(self, values: Sequence[float]) -> None:
|
||||
"""A wrong-length vector is a corpus-wide defect; fail on the first."""
|
||||
if len(values) != self.dimensions:
|
||||
raise ValueError(
|
||||
f"{self.model_id} returned {len(values)} dimensions, "
|
||||
f"expected {self.dimensions}"
|
||||
)
|
||||
@@ -0,0 +1,71 @@
|
||||
"""One live call to one provider — the smallest thing that proves it works.
|
||||
|
||||
Deliberately not wired into `ingestion.cli`: that module is being edited for
|
||||
the parser/chunking work, and this task has no business touching it. Run as
|
||||
|
||||
python -m ingestion.embed.probe --provider titan-v2
|
||||
|
||||
Each run embeds a **single short string** and makes exactly one request, so it
|
||||
answers "are the credentials, the model access and my request shape all
|
||||
right?" without approaching the cost or the risk of a corpus run. It reports
|
||||
the request key shape, the returned dimension, the measured L2 norm and the
|
||||
latency — the L2 norm because whether a provider returns unit-length vectors
|
||||
is a property worth measuring rather than reading off a documentation page.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import math
|
||||
import sys
|
||||
from typing import Sequence
|
||||
|
||||
from .ports import INPUT_DOCUMENT, INPUT_KINDS, EmbeddingVector
|
||||
from .registry import DEFAULT_REGION, build_provider, provider_names
|
||||
|
||||
# Vietnamese, diacritics, and a dose string: the shape of the real corpus, not
|
||||
# "hello world". If an encoding path is broken this is what shows it.
|
||||
DEFAULT_TEXT = "Liều dùng: uống 500 mg mỗi 8 giờ, không quá 4 g mỗi ngày."
|
||||
|
||||
|
||||
def _l2_norm(values: Sequence[float]) -> float:
|
||||
return math.sqrt(sum(v * v for v in values))
|
||||
|
||||
|
||||
def _report(vector: EmbeddingVector, latency_ms: float, requests: int) -> None:
|
||||
print(f"provider : {vector.provider}")
|
||||
print(f"model_id : {vector.model_id}")
|
||||
print(f"input_kind : {vector.input_kind}")
|
||||
print(f"requests : {requests}")
|
||||
print(f"dimensions : {len(vector.values)} (expected {vector.dimensions})")
|
||||
print(f"measured L2 norm: {_l2_norm(vector.values):.6f}")
|
||||
print(f"normalized (per provider docs): {vector.normalized}")
|
||||
print(f"input_token_count: {vector.input_token_count}")
|
||||
print(f"latency_ms : {latency_ms:.1f}")
|
||||
print(f"first 5 values : {[round(v, 6) for v in vector.values[:5]]}")
|
||||
|
||||
|
||||
def main(argv=None) -> int:
|
||||
parser = argparse.ArgumentParser(
|
||||
prog="python -m ingestion.embed.probe",
|
||||
description="Make ONE live embedding call and report what came back.",
|
||||
)
|
||||
parser.add_argument("--provider", required=True, choices=provider_names())
|
||||
parser.add_argument("--region", default=DEFAULT_REGION)
|
||||
parser.add_argument("--dimensions", type=int, default=None)
|
||||
parser.add_argument("--text", default=DEFAULT_TEXT)
|
||||
parser.add_argument("--input-kind", default=INPUT_DOCUMENT, choices=INPUT_KINDS)
|
||||
args = parser.parse_args(argv)
|
||||
|
||||
if hasattr(sys.stdout, "reconfigure"):
|
||||
sys.stdout.reconfigure(encoding="utf-8")
|
||||
|
||||
provider = build_provider(
|
||||
args.provider, region=args.region, dimensions=args.dimensions
|
||||
)
|
||||
batch = provider.embed([args.text], args.input_kind)
|
||||
_report(batch.vectors[0], batch.latency_ms, batch.request_count)
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@@ -0,0 +1,78 @@
|
||||
"""Name -> provider, so nothing downstream has to import an SDK to pick one.
|
||||
|
||||
Adding a fourth model to the benchmark means adding one entry to `_BUILDERS`,
|
||||
not editing a caller — the open/closed rule CLAUDE.md applies to the section
|
||||
taxonomy, applied here for the same reason.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Callable, Dict, Optional
|
||||
|
||||
from . import bedrock_cohere, bedrock_titan, local_bge_m3
|
||||
from .bedrock_runtime import DEFAULT_REGION, BedrockInvoker, Boto3BedrockInvoker
|
||||
from .ports import EmbeddingProvider
|
||||
|
||||
TITAN_V2 = bedrock_titan.PROVIDER_NAME
|
||||
COHERE_V4 = bedrock_cohere.PROVIDER_NAME
|
||||
BGE_M3 = local_bge_m3.PROVIDER_NAME
|
||||
|
||||
CLOUD_PROVIDERS = (TITAN_V2, COHERE_V4)
|
||||
|
||||
|
||||
def _build_titan(
|
||||
invoker: Optional[BedrockInvoker], region: str, dimensions: Optional[int]
|
||||
) -> EmbeddingProvider:
|
||||
return bedrock_titan.TitanTextEmbeddingsV2(
|
||||
invoker or Boto3BedrockInvoker(region=region),
|
||||
dimensions=dimensions or 1024,
|
||||
)
|
||||
|
||||
|
||||
def _build_cohere(
|
||||
invoker: Optional[BedrockInvoker], region: str, dimensions: Optional[int]
|
||||
) -> EmbeddingProvider:
|
||||
return bedrock_cohere.CohereEmbedV4(
|
||||
invoker or Boto3BedrockInvoker(region=region),
|
||||
dimensions=dimensions or 1024,
|
||||
)
|
||||
|
||||
|
||||
def _build_bge_m3(
|
||||
_invoker: Optional[BedrockInvoker], _region: str, dimensions: Optional[int]
|
||||
) -> EmbeddingProvider:
|
||||
# Runs on this machine: there is no invoker and no region to honour.
|
||||
if dimensions not in (None, local_bge_m3.DENSE_DIMENSIONS):
|
||||
raise ValueError(
|
||||
f"{BGE_M3} produces {local_bge_m3.DENSE_DIMENSIONS} dimensions; "
|
||||
f"{dimensions} was requested"
|
||||
)
|
||||
return local_bge_m3.BgeM3Local()
|
||||
|
||||
|
||||
Builder = Callable[[Optional[BedrockInvoker], str, Optional[int]], EmbeddingProvider]
|
||||
|
||||
_BUILDERS: Dict[str, Builder] = {
|
||||
TITAN_V2: _build_titan,
|
||||
COHERE_V4: _build_cohere,
|
||||
BGE_M3: _build_bge_m3,
|
||||
}
|
||||
|
||||
|
||||
def provider_names() -> tuple:
|
||||
return tuple(_BUILDERS)
|
||||
|
||||
|
||||
def build_provider(
|
||||
name: str,
|
||||
*,
|
||||
invoker: Optional[BedrockInvoker] = None,
|
||||
region: str = DEFAULT_REGION,
|
||||
dimensions: Optional[int] = None,
|
||||
) -> EmbeddingProvider:
|
||||
try:
|
||||
builder = _BUILDERS[name]
|
||||
except KeyError:
|
||||
raise ValueError(
|
||||
f"unknown embedding provider {name!r}; known: {provider_names()}"
|
||||
) from None
|
||||
return builder(invoker, region, dimensions)
|
||||
@@ -0,0 +1 @@
|
||||
"""Verified drug-entity artifact builders."""
|
||||
@@ -0,0 +1,156 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import re
|
||||
import unicodedata
|
||||
from difflib import SequenceMatcher
|
||||
from pathlib import Path
|
||||
|
||||
import fitz
|
||||
|
||||
from ingestion.validation.back_index import parse_back_index_see_aliases
|
||||
|
||||
WORD_RE = re.compile(r"\w+", re.UNICODE)
|
||||
PAREN_RE = re.compile(r"\(([^()]*)\)")
|
||||
|
||||
|
||||
def normalize_name(text: str) -> str:
|
||||
decomposed = unicodedata.normalize("NFKD", text.casefold()).replace("đ", "d")
|
||||
plain = "".join(char for char in decomposed if not unicodedata.combining(char))
|
||||
return " ".join(WORD_RE.findall(plain))
|
||||
|
||||
|
||||
def _canonical_aliases(drug_name: str) -> set[str]:
|
||||
aliases = {drug_name.strip()}
|
||||
without_parentheses = PAREN_RE.sub("", drug_name).strip()
|
||||
if without_parentheses:
|
||||
aliases.add(without_parentheses)
|
||||
aliases.update(
|
||||
value.strip() for value in PAREN_RE.findall(drug_name) if value.strip()
|
||||
)
|
||||
return aliases
|
||||
|
||||
|
||||
def _trade_aliases(monograph: dict) -> set[str]:
|
||||
section = monograph.get("sections", {}).get("ten_thuong_mai")
|
||||
if not section:
|
||||
return set()
|
||||
text = section.get("text", "")
|
||||
return {
|
||||
value.strip().strip(".")
|
||||
for value in re.split(r"[,;\n]", text)
|
||||
if value.strip().strip(".")
|
||||
}
|
||||
|
||||
|
||||
def _read_monographs(path: Path) -> list[dict]:
|
||||
with path.open(encoding="utf-8") as handle:
|
||||
return [json.loads(line) for line in handle if line.strip()]
|
||||
|
||||
|
||||
def build_entities(monographs_path: Path, pdf_path: Path) -> dict:
|
||||
monographs = _read_monographs(monographs_path)
|
||||
entities: dict[str, dict] = {}
|
||||
lookup: list[tuple[str, str, tuple[int, int]]] = []
|
||||
for monograph in monographs:
|
||||
drug_id = monograph["drug_id"]
|
||||
aliases = _canonical_aliases(monograph["drug_name"])
|
||||
aliases.update(_trade_aliases(monograph))
|
||||
aliases.add(drug_id.replace("_", " "))
|
||||
entities[drug_id] = {
|
||||
"drug_id": drug_id,
|
||||
"canonical_name": monograph["drug_name"],
|
||||
"aliases": aliases,
|
||||
"atc_codes": sorted(set(monograph.get("atc_codes", []))),
|
||||
"source_page_range": monograph.get("source_page_range"),
|
||||
}
|
||||
page_range = tuple(monograph["source_page_range"])
|
||||
lookup.extend(
|
||||
(normalize_name(alias), drug_id, page_range) for alias in _canonical_aliases(
|
||||
monograph["drug_name"],
|
||||
) if normalize_name(alias)
|
||||
)
|
||||
|
||||
unresolved = []
|
||||
ambiguous = []
|
||||
with fitz.open(pdf_path) as doc:
|
||||
index_aliases = parse_back_index_see_aliases(doc)
|
||||
for relation in index_aliases:
|
||||
target = normalize_name(relation.target)
|
||||
candidates = {
|
||||
drug_id for alias, drug_id, page_range in lookup
|
||||
if page_range[0] <= relation.printed_page - 1 <= page_range[1]
|
||||
and (
|
||||
target == alias
|
||||
or target.startswith(f"{alias} ")
|
||||
or target.endswith(f" {alias}")
|
||||
or f" {alias} " in target
|
||||
)
|
||||
}
|
||||
if not candidates:
|
||||
fuzzy = sorted(
|
||||
(
|
||||
SequenceMatcher(None, target, alias).ratio(),
|
||||
drug_id,
|
||||
)
|
||||
for alias, drug_id, page_range in lookup
|
||||
if page_range[0] <= relation.printed_page - 1 <= page_range[1]
|
||||
)
|
||||
if fuzzy and fuzzy[-1][0] >= 0.9:
|
||||
runner_up = fuzzy[-2][0] if len(fuzzy) > 1 else 0.0
|
||||
if fuzzy[-1][0] - runner_up >= 0.05:
|
||||
candidates = {fuzzy[-1][1]}
|
||||
if len(candidates) == 1:
|
||||
entities[next(iter(candidates))]["aliases"].add(relation.alias)
|
||||
elif candidates:
|
||||
ambiguous.append(relation.alias)
|
||||
else:
|
||||
unresolved.append(relation.alias)
|
||||
|
||||
output_entities = []
|
||||
for entity in entities.values():
|
||||
output_entities.append({
|
||||
**entity,
|
||||
"aliases": sorted(entity["aliases"], key=lambda value: value.casefold()),
|
||||
})
|
||||
return {
|
||||
"schema_version": 1,
|
||||
"entities": sorted(output_entities, key=lambda item: item["drug_id"]),
|
||||
"stats": {
|
||||
"entity_count": len(output_entities),
|
||||
"back_index_see_relations": len(index_aliases),
|
||||
"back_index_aliases_mapped": len(index_aliases) - len(unresolved) - len(ambiguous),
|
||||
"back_index_aliases_unresolved": len(unresolved),
|
||||
"back_index_aliases_ambiguous": len(ambiguous),
|
||||
"trade_name_sections": sum(
|
||||
"ten_thuong_mai" in row.get("sections", {}) for row in monographs
|
||||
),
|
||||
"total_aliases": sum(len(row["aliases"]) for row in output_entities),
|
||||
},
|
||||
"unresolved_back_index_aliases": sorted(unresolved),
|
||||
"ambiguous_back_index_aliases": sorted(ambiguous),
|
||||
}
|
||||
|
||||
|
||||
def write_entities(payload: dict, path: Path) -> None:
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
path.write_text(
|
||||
json.dumps(payload, ensure_ascii=False, indent=2) + "\n",
|
||||
encoding="utf-8",
|
||||
)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--monographs", type=Path, required=True)
|
||||
parser.add_argument("--pdf", type=Path, required=True)
|
||||
parser.add_argument("--output", type=Path, required=True)
|
||||
args = parser.parse_args()
|
||||
payload = build_entities(args.monographs, args.pdf)
|
||||
write_entities(payload, args.output)
|
||||
print(json.dumps(payload["stats"], ensure_ascii=False))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -26,6 +26,12 @@ from ..tables.models import TableRegion
|
||||
# One line of body type on this book measures ~10.5pt; a fraction spans the
|
||||
# numerator line, the bar and the denominator line.
|
||||
FORMULA_BAND_HEIGHT_PT = 13.0
|
||||
# ADENOSIN p147 is printed as numerator and denominator lines with no rule.
|
||||
# Its synthetic anchor sits on the wrapped numerator baseline, so symmetric
|
||||
# growth covers only the numerator. The denominator ends ~29pt below the
|
||||
# anchor; the next "Ví dụ:" line has its centre ~32.3pt below. 31pt captures
|
||||
# the denominator while excluding that prose and the following table.
|
||||
BARLESS_FORMULA_BOTTOM_PT = 31.0
|
||||
|
||||
# Wide on purpose. The bar is often narrower than the numerator above it, and
|
||||
# a numerator span can carry leading spaces that push its box's centre well to
|
||||
@@ -51,12 +57,16 @@ def load_formula_regions(path: Path | None = None) -> List[TableRegion]:
|
||||
regions = []
|
||||
for index, entry in enumerate(payload["regions"]):
|
||||
x0, y0, x1, y1 = entry["bar_bbox"]
|
||||
bottom_margin = (
|
||||
BARLESS_FORMULA_BOTTOM_PT
|
||||
if entry.get("source_prints_no_bar") else FORMULA_BAND_HEIGHT_PT
|
||||
)
|
||||
regions.append(
|
||||
TableRegion(
|
||||
table_id=f"p{entry['physical_page']}_f{index}",
|
||||
physical_page=entry["physical_page"],
|
||||
bbox=(x0 - FORMULA_SIDE_MARGIN_PT, y0 - FORMULA_BAND_HEIGHT_PT,
|
||||
x1 + FORMULA_SIDE_MARGIN_PT, y1 + FORMULA_BAND_HEIGHT_PT),
|
||||
x1 + FORMULA_SIDE_MARGIN_PT, y1 + bottom_margin),
|
||||
n_rows=2,
|
||||
n_cols=1,
|
||||
shape=SHAPE_FORMULA_2D,
|
||||
|
||||
@@ -32,6 +32,10 @@ class Span:
|
||||
def bold(self) -> bool:
|
||||
return "Bold" in self.font
|
||||
|
||||
@property
|
||||
def italic(self) -> bool:
|
||||
return "Italic" in self.font or "Oblique" in self.font
|
||||
|
||||
@property
|
||||
def span_id(self) -> str:
|
||||
"""Stable identifier for one source span.
|
||||
|
||||
@@ -0,0 +1,68 @@
|
||||
"""Load stage: chunks plus vectors into a vector store, idempotently.
|
||||
|
||||
Import from here rather than from a module. Nothing outside this package should
|
||||
name `qdrant_client` — that dependency is confined to `qdrant_repo`, so the
|
||||
whole stage runs and is tested with no server, exactly as `embed/` confines
|
||||
boto3 to `bedrock_runtime`.
|
||||
|
||||
The two rules this package exists to enforce are worth naming at the front
|
||||
door. A point id is derived from `chunk_id`, so loading twice converges instead
|
||||
of duplicating. And a collection records the corpus digest and model it was
|
||||
built from, so a second corpus or a second model cannot be mixed into it — a
|
||||
state that produces no error at query time and would otherwise be invisible.
|
||||
"""
|
||||
from .corpus import corpus_sha256, count_chunks, iter_chunk_records
|
||||
from .in_memory import InMemoryVectorStore
|
||||
from .manifest import (
|
||||
MANIFEST_SUFFIX,
|
||||
CorpusMismatch,
|
||||
assert_compatible,
|
||||
manifest_collection,
|
||||
read_manifest,
|
||||
write_manifest,
|
||||
)
|
||||
from .models import (
|
||||
COSINE,
|
||||
INDEXED_PAYLOAD_FIELDS,
|
||||
REQUIRED_CHUNK_FIELDS,
|
||||
CollectionSpec,
|
||||
CorpusManifest,
|
||||
VectorPoint,
|
||||
build_point,
|
||||
point_id_for,
|
||||
validate_chunk_record,
|
||||
)
|
||||
from .ports import VectorStore
|
||||
from .upsert import (
|
||||
DEFAULT_BATCH_SIZE,
|
||||
ChunkLoader,
|
||||
LoadReport,
|
||||
PointCountMismatch,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"COSINE",
|
||||
"DEFAULT_BATCH_SIZE",
|
||||
"INDEXED_PAYLOAD_FIELDS",
|
||||
"MANIFEST_SUFFIX",
|
||||
"REQUIRED_CHUNK_FIELDS",
|
||||
"ChunkLoader",
|
||||
"CollectionSpec",
|
||||
"CorpusManifest",
|
||||
"CorpusMismatch",
|
||||
"InMemoryVectorStore",
|
||||
"LoadReport",
|
||||
"PointCountMismatch",
|
||||
"VectorPoint",
|
||||
"VectorStore",
|
||||
"assert_compatible",
|
||||
"build_point",
|
||||
"corpus_sha256",
|
||||
"count_chunks",
|
||||
"iter_chunk_records",
|
||||
"manifest_collection",
|
||||
"point_id_for",
|
||||
"read_manifest",
|
||||
"validate_chunk_record",
|
||||
"write_manifest",
|
||||
]
|
||||
|
||||
@@ -0,0 +1,48 @@
|
||||
"""Reading the chunk artifact, and computing the identity it is loaded under.
|
||||
|
||||
The digest is taken over content rather than over parsed records, so it still
|
||||
changes on a field reordering that leaves every record semantically identical —
|
||||
a normalised, semantic digest would let a regenerated corpus pass the A6 gate
|
||||
while its point payloads no longer match what `chunk/` produces.
|
||||
|
||||
**Line endings are the one thing normalised, and only because not doing it was
|
||||
a bug.** A raw-byte digest makes the same JSONL hash differently after a Windows
|
||||
checkout with CRLF than after a Linux one with LF. The A6 gate would then refuse
|
||||
a load in CI against a corpus that is byte-for-byte the same data, which is a
|
||||
false rejection with a confusing message — the failure mode of a safety gate
|
||||
that cries wolf is that someone turns it off. Each line is digested with its
|
||||
terminator normalised to `\\n`, which costs none of the strictness that matters.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, Iterator
|
||||
|
||||
def corpus_sha256(path: Path | str) -> str:
|
||||
"""Content digest, identical across CRLF and LF checkouts of the same data."""
|
||||
digest = hashlib.sha256()
|
||||
with Path(path).open("rb") as handle:
|
||||
for raw in handle:
|
||||
digest.update(raw.rstrip(b"\r\n"))
|
||||
digest.update(b"\n")
|
||||
return digest.hexdigest()
|
||||
|
||||
|
||||
def iter_chunk_records(path: Path | str) -> Iterator[Dict[str, Any]]:
|
||||
with Path(path).open("r", encoding="utf-8") as handle:
|
||||
for line_number, line in enumerate(handle, start=1):
|
||||
stripped = line.strip()
|
||||
if not stripped:
|
||||
continue
|
||||
try:
|
||||
yield json.loads(stripped)
|
||||
except json.JSONDecodeError as exc:
|
||||
raise ValueError(
|
||||
f"{Path(path).name} line {line_number} is not valid JSON: {exc}"
|
||||
) from exc
|
||||
|
||||
|
||||
def count_chunks(path: Path | str) -> int:
|
||||
return sum(1 for _ in iter_chunk_records(path))
|
||||
@@ -0,0 +1,93 @@
|
||||
"""A `VectorStore` that keeps everything in a dict.
|
||||
|
||||
Not only a test double. It is the reference implementation of the port: the
|
||||
idempotency rule (same point id overwrites, never appends) and the
|
||||
unknown-collection failures are stated here once, so a test that passes against
|
||||
this store is testing the contract rather than an accident of Qdrant's
|
||||
behaviour. `apps/ai-service/rag` already keeps an in-memory store for the same
|
||||
reason.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, Dict, List, Mapping, Optional, Sequence
|
||||
|
||||
from .models import CollectionSpec, VectorPoint
|
||||
|
||||
|
||||
class InMemoryVectorStore:
|
||||
def __init__(self) -> None:
|
||||
self._collections: Dict[str, CollectionSpec] = {}
|
||||
self._points: Dict[str, Dict[str, VectorPoint]] = {}
|
||||
self._indexes: Dict[str, List[tuple]] = {}
|
||||
|
||||
def collection_exists(self, name: str) -> bool:
|
||||
return name in self._collections
|
||||
|
||||
def create_collection(self, spec: CollectionSpec) -> None:
|
||||
if spec.name in self._collections:
|
||||
raise ValueError(f"collection {spec.name!r} already exists")
|
||||
self._collections[spec.name] = spec
|
||||
self._points[spec.name] = {}
|
||||
self._indexes[spec.name] = []
|
||||
|
||||
def delete_collection(self, name: str) -> None:
|
||||
self._collections.pop(name, None)
|
||||
self._points.pop(name, None)
|
||||
self._indexes.pop(name, None)
|
||||
|
||||
def create_payload_index(
|
||||
self, name: str, field_name: str, field_schema: str
|
||||
) -> None:
|
||||
self._require(name)
|
||||
self._indexes[name].append((field_name, field_schema))
|
||||
|
||||
def upsert(self, name: str, points: Sequence[VectorPoint]) -> int:
|
||||
self._require(name)
|
||||
spec = self._collections[name]
|
||||
for point in points:
|
||||
if len(point.vector) != spec.vector_size:
|
||||
raise ValueError(
|
||||
f"point {point.id} has {len(point.vector)} dimensions, "
|
||||
f"collection {name!r} expects {spec.vector_size}"
|
||||
)
|
||||
self._points[name][point.id] = point
|
||||
return len(points)
|
||||
|
||||
def count(self, name: str) -> int:
|
||||
self._require(name)
|
||||
return len(self._points[name])
|
||||
|
||||
def retrieve(self, name: str, point_id: str) -> Optional[VectorPoint]:
|
||||
self._require(name)
|
||||
return self._points[name].get(point_id)
|
||||
|
||||
def find_by_payload(
|
||||
self, name: str, equals: Mapping[str, Any]
|
||||
) -> List[VectorPoint]:
|
||||
self._require(name)
|
||||
if not equals:
|
||||
raise ValueError("find_by_payload needs at least one condition")
|
||||
return [
|
||||
point
|
||||
for point in self._points[name].values()
|
||||
if all(_matches(point.payload.get(k), v) for k, v in equals.items())
|
||||
]
|
||||
|
||||
def indexed_fields(self, name: str) -> List[tuple]:
|
||||
self._require(name)
|
||||
return list(self._indexes[name])
|
||||
|
||||
def spec(self, name: str) -> CollectionSpec:
|
||||
self._require(name)
|
||||
return self._collections[name]
|
||||
|
||||
def _require(self, name: str) -> None:
|
||||
if name not in self._collections:
|
||||
raise KeyError(f"collection {name!r} does not exist")
|
||||
|
||||
|
||||
def _matches(stored: Any, wanted: Any) -> bool:
|
||||
"""Qdrant's rule: a list field matches when any element equals the value."""
|
||||
if isinstance(stored, list):
|
||||
return wanted in stored
|
||||
return stored == wanted
|
||||
@@ -0,0 +1,86 @@
|
||||
"""Binding a collection to the exact corpus and model it was built from (A6).
|
||||
|
||||
**Why a sidecar collection rather than a reserved point.** Qdrant has no
|
||||
collection-level metadata field, so the manifest has to live in a point. Putting
|
||||
that point inside the data collection would make `count()` one larger than the
|
||||
chunk count — and `qdrant_point_count != chunk_count` is a v1 acceptance gate
|
||||
(`docs/v1-delivery-plan.md` §6). A gate that needs an "except the manifest"
|
||||
footnote is a gate that will eventually be read wrong. A `<name>__manifest`
|
||||
collection keeps the data collection's count exactly equal to the number of
|
||||
chunks, and keeps the manifest out of every search result by construction
|
||||
rather than by remembering to filter it.
|
||||
|
||||
The vector on that point is a single zero. It is never searched; the point
|
||||
exists only to carry a payload.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Optional
|
||||
|
||||
from .models import CollectionSpec, CorpusManifest, VectorPoint
|
||||
from .ports import VectorStore
|
||||
|
||||
MANIFEST_SUFFIX = "__manifest"
|
||||
MANIFEST_POINT_ID = "00000000-0000-5000-8000-000000000001"
|
||||
|
||||
|
||||
class CorpusMismatch(RuntimeError):
|
||||
"""Raised instead of upserting a corpus into a collection built elsewhere."""
|
||||
|
||||
|
||||
def manifest_collection(name: str) -> str:
|
||||
return f"{name}{MANIFEST_SUFFIX}"
|
||||
|
||||
|
||||
def write_manifest(store: VectorStore, name: str, manifest: CorpusManifest) -> None:
|
||||
sidecar = manifest_collection(name)
|
||||
if not store.collection_exists(sidecar):
|
||||
store.create_collection(CollectionSpec(name=sidecar, vector_size=1))
|
||||
store.upsert(
|
||||
sidecar,
|
||||
[
|
||||
VectorPoint(
|
||||
id=MANIFEST_POINT_ID,
|
||||
vector=[0.0],
|
||||
payload=manifest.to_payload(),
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
def read_manifest(store: VectorStore, name: str) -> Optional[CorpusManifest]:
|
||||
sidecar = manifest_collection(name)
|
||||
if not store.collection_exists(sidecar):
|
||||
return None
|
||||
point = store.retrieve(sidecar, MANIFEST_POINT_ID)
|
||||
if point is None:
|
||||
return None
|
||||
return CorpusManifest.from_payload(point.payload)
|
||||
|
||||
|
||||
def assert_compatible(
|
||||
store: VectorStore, name: str, incoming: CorpusManifest
|
||||
) -> Optional[CorpusManifest]:
|
||||
"""Refuse the load unless the collection was built from the same corpus.
|
||||
|
||||
Returns the stored manifest, or None when the collection is new. An
|
||||
existing data collection with no manifest is itself a refusal: it was
|
||||
loaded by something that did not record what it loaded, so nothing can be
|
||||
said about what is already in there.
|
||||
"""
|
||||
stored = read_manifest(store, name)
|
||||
if stored is None:
|
||||
if store.collection_exists(name) and store.count(name) > 0:
|
||||
raise CorpusMismatch(
|
||||
f"collection {name!r} already holds {store.count(name)} points but "
|
||||
f"has no manifest; refusing to mix an unknown corpus with "
|
||||
f"{incoming.corpus_sha256[:12]}"
|
||||
)
|
||||
return None
|
||||
|
||||
reasons = stored.conflicts_with(incoming)
|
||||
if reasons:
|
||||
raise CorpusMismatch(
|
||||
f"refusing to load into {name!r}: " + "; ".join(reasons)
|
||||
)
|
||||
return stored
|
||||
@@ -0,0 +1,243 @@
|
||||
"""What goes into the vector store, and the rules a record must satisfy first.
|
||||
|
||||
Two decisions are worth stating because their alternatives look reasonable.
|
||||
|
||||
**Point ids are derived, never generated.** `uuid5` of `chunk_id` means the
|
||||
same chunk always lands on the same point, so a re-run overwrites rather than
|
||||
duplicates. A random id would make `cli load` non-idempotent, and the damage
|
||||
would be invisible — the collection would simply hold two copies of a dose and
|
||||
return whichever ranked higher.
|
||||
|
||||
**Payload carries the whole chunk record, not a chosen subset.** CLAUDE.md
|
||||
requires provenance to survive a stage boundary; a whitelist here would silently
|
||||
drop any field a later chunker adds, which is exactly the failure it warns
|
||||
about. So the record passes through intact and only a required core is
|
||||
*checked* — `population_tags` or `printed_page_range` will flow through the day
|
||||
`chunk/` starts emitting them, with no edit to this module.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import uuid
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Dict, List, Mapping, Optional, Sequence
|
||||
|
||||
# Fixed namespace: changing it would re-id the entire corpus and orphan every
|
||||
# point already loaded. It is a constant of the project, not a tunable.
|
||||
POINT_NAMESPACE = uuid.UUID("6f0d6d1e-4c2a-5f6b-9a3d-2f8e1c7b4a90")
|
||||
|
||||
COSINE = "Cosine"
|
||||
|
||||
# Schema v4 separates retrieval-enriched `text` from contiguous `source_text`.
|
||||
# Clinicians cite the printed folio, and
|
||||
# `citation_uses_physical_page = 0` is a v1 acceptance gate, so a chunk without
|
||||
# one cannot be cited honestly. Refusing it here is the difference between
|
||||
# noticing before an embedding run and noticing after paying for one, when every
|
||||
# answer abstains for missing provenance.
|
||||
SUPPORTED_SCHEMA_VERSION = 4
|
||||
|
||||
# Absence means the loader was handed something other than the chunk artifact,
|
||||
# and it should say so loudly rather than write a point that cannot be traced
|
||||
# back to a page.
|
||||
REQUIRED_CHUNK_FIELDS = (
|
||||
"chunk_id",
|
||||
"drug_id",
|
||||
"drug_name",
|
||||
"section_key",
|
||||
"text",
|
||||
"source_text",
|
||||
"heading_physical_page",
|
||||
"source_page_range",
|
||||
"printed_page_range",
|
||||
"chunk_kind",
|
||||
)
|
||||
|
||||
# Both must be a real two-page span, not merely present. An empty list is the
|
||||
# shape a defaulted field takes, and `[] in (None, "")` is False — which is how
|
||||
# an earlier version of this check passed one.
|
||||
PAGE_RANGE_FIELDS = ("source_page_range", "printed_page_range")
|
||||
|
||||
# Fields the retrieval side filters on. Qdrant needs an explicit index per
|
||||
# field; without it a filtered query still works but scans.
|
||||
INDEXED_PAYLOAD_FIELDS = (
|
||||
("chunk_id", "keyword"),
|
||||
("drug_id", "keyword"),
|
||||
("section_key", "keyword"),
|
||||
("atc_codes", "keyword"),
|
||||
("chunk_kind", "keyword"),
|
||||
("has_quarantined_content", "bool"),
|
||||
)
|
||||
|
||||
|
||||
def point_id_for(chunk_id: str) -> str:
|
||||
if not chunk_id or not chunk_id.strip():
|
||||
raise ValueError("chunk_id is required to derive a point id")
|
||||
return str(uuid.uuid5(POINT_NAMESPACE, chunk_id))
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class VectorPoint:
|
||||
id: str
|
||||
vector: List[float]
|
||||
payload: Dict[str, Any]
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class CollectionSpec:
|
||||
name: str
|
||||
vector_size: int
|
||||
distance: str = COSINE
|
||||
indexed_fields: Sequence[tuple] = INDEXED_PAYLOAD_FIELDS
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
if not self.name.strip():
|
||||
raise ValueError("collection name is required")
|
||||
if self.vector_size <= 0:
|
||||
raise ValueError(f"vector_size must be positive, got {self.vector_size}")
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class CorpusManifest:
|
||||
"""What a collection was built from — the thing A6 compares against.
|
||||
|
||||
Mixing two generations of the corpus, or two models' vectors, into one
|
||||
collection produces no error at query time: every search still returns
|
||||
something. This record is what makes that state detectable instead.
|
||||
"""
|
||||
|
||||
corpus_sha256: str
|
||||
chunk_count: int
|
||||
model_id: str
|
||||
dimensions: int
|
||||
input_kind: str
|
||||
provider: Optional[str] = None
|
||||
distance: str = COSINE
|
||||
extras: Dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
def to_payload(self) -> Dict[str, Any]:
|
||||
return {
|
||||
"corpus_sha256": self.corpus_sha256,
|
||||
"chunk_count": self.chunk_count,
|
||||
"model_id": self.model_id,
|
||||
"dimensions": self.dimensions,
|
||||
"input_kind": self.input_kind,
|
||||
"provider": self.provider,
|
||||
"distance": self.distance,
|
||||
**self.extras,
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def from_payload(cls, payload: Mapping[str, Any]) -> "CorpusManifest":
|
||||
known = {
|
||||
"corpus_sha256",
|
||||
"chunk_count",
|
||||
"model_id",
|
||||
"dimensions",
|
||||
"input_kind",
|
||||
"provider",
|
||||
"distance",
|
||||
}
|
||||
return cls(
|
||||
corpus_sha256=payload["corpus_sha256"],
|
||||
chunk_count=payload["chunk_count"],
|
||||
model_id=payload["model_id"],
|
||||
dimensions=payload["dimensions"],
|
||||
input_kind=payload["input_kind"],
|
||||
provider=payload.get("provider"),
|
||||
distance=payload.get("distance", COSINE),
|
||||
extras={k: v for k, v in payload.items() if k not in known},
|
||||
)
|
||||
|
||||
def conflicts_with(self, other: "CorpusManifest") -> List[str]:
|
||||
"""Every reason these two must not share a collection."""
|
||||
reasons = []
|
||||
if self.corpus_sha256 != other.corpus_sha256:
|
||||
reasons.append(
|
||||
f"corpus sha256 {other.corpus_sha256[:12]} does not match the "
|
||||
f"{self.corpus_sha256[:12]} this collection was built from"
|
||||
)
|
||||
if self.model_id != other.model_id:
|
||||
reasons.append(
|
||||
f"model {other.model_id} does not match the collection's "
|
||||
f"{self.model_id}"
|
||||
)
|
||||
if self.dimensions != other.dimensions:
|
||||
reasons.append(
|
||||
f"{other.dimensions} dimensions do not match the collection's "
|
||||
f"{self.dimensions}"
|
||||
)
|
||||
if self.input_kind != other.input_kind:
|
||||
reasons.append(
|
||||
f"input kind {other.input_kind} does not match the collection's "
|
||||
f"{self.input_kind}"
|
||||
)
|
||||
return reasons
|
||||
|
||||
|
||||
def _is_missing(value: Any) -> bool:
|
||||
"""`0` and `False` are values; an empty string or empty list is not.
|
||||
|
||||
Physical page 0 and `has_quarantined_content=False` are both legitimate, so
|
||||
a plain falsiness test would reject real records. Only `None` and empty
|
||||
collections count as absent.
|
||||
"""
|
||||
if value is None:
|
||||
return True
|
||||
return isinstance(value, (str, list, tuple, dict, set)) and len(value) == 0
|
||||
|
||||
|
||||
def validate_chunk_record(record: Mapping[str, Any]) -> None:
|
||||
label = record.get("chunk_id", "<no chunk_id>")
|
||||
|
||||
missing = [f for f in REQUIRED_CHUNK_FIELDS if _is_missing(record.get(f))]
|
||||
if missing:
|
||||
raise ValueError(
|
||||
f"chunk record {label!r} is missing "
|
||||
f"required provenance fields: {missing}"
|
||||
)
|
||||
|
||||
version = record.get("schema_version")
|
||||
if type(version) is not int or version != SUPPORTED_SCHEMA_VERSION:
|
||||
raise ValueError(
|
||||
f"chunk record {label!r} declares schema_version {version!r}; "
|
||||
f"this loader supports exactly v{SUPPORTED_SCHEMA_VERSION}; "
|
||||
"unknown older or newer schemas are refused until compatibility "
|
||||
"is implemented explicitly"
|
||||
)
|
||||
|
||||
for field_name in PAGE_RANGE_FIELDS:
|
||||
_validate_page_range(label, field_name, record.get(field_name))
|
||||
|
||||
|
||||
def _validate_page_range(label: str, field_name: str, value: Any) -> None:
|
||||
if not isinstance(value, (list, tuple)) or len(value) != 2:
|
||||
raise ValueError(
|
||||
f"chunk record {label!r} has {field_name}={value!r}; expected a "
|
||||
f"[start, end] pair"
|
||||
)
|
||||
start, end = value
|
||||
if type(start) is not int or type(end) is not int:
|
||||
raise ValueError(
|
||||
f"chunk record {label!r} has a non-integer page in "
|
||||
f"{field_name}={value!r}"
|
||||
)
|
||||
if start > end:
|
||||
raise ValueError(
|
||||
f"chunk record {label!r} has {field_name}={value!r} running backwards"
|
||||
)
|
||||
minimum = 1 if field_name == "printed_page_range" else 0
|
||||
if start < minimum:
|
||||
raise ValueError(
|
||||
f"chunk record {label!r} has {field_name}={value!r}; "
|
||||
f"pages must start at {minimum} or later"
|
||||
)
|
||||
|
||||
|
||||
def build_point(
|
||||
record: Mapping[str, Any], vector: Sequence[float]
|
||||
) -> VectorPoint:
|
||||
validate_chunk_record(record)
|
||||
return VectorPoint(
|
||||
id=point_id_for(record["chunk_id"]),
|
||||
vector=list(vector),
|
||||
payload=dict(record),
|
||||
)
|
||||
@@ -0,0 +1,52 @@
|
||||
"""The vector-store boundary.
|
||||
|
||||
The loader, the corpus gate and every test in this package talk to this
|
||||
Protocol and never to a database SDK. That is what lets the whole load stage be
|
||||
exercised offline against `InMemoryVectorStore`, and what keeps `qdrant_client`
|
||||
named in exactly one module — the same arrangement `embed/` uses to confine
|
||||
boto3 to `bedrock_runtime`.
|
||||
|
||||
Deliberately primitive: collections, points, payload indexes. Anything with an
|
||||
opinion — how a corpus is bound to a collection, how a chunk becomes a point —
|
||||
is a layer above this, so swapping the store does not drag those rules with it.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, Mapping, Optional, Protocol, Sequence
|
||||
|
||||
from .models import CollectionSpec, VectorPoint
|
||||
|
||||
|
||||
class VectorStore(Protocol):
|
||||
def collection_exists(self, name: str) -> bool: ...
|
||||
|
||||
def create_collection(self, spec: CollectionSpec) -> None: ...
|
||||
|
||||
def delete_collection(self, name: str) -> None: ...
|
||||
|
||||
def create_payload_index(
|
||||
self, name: str, field_name: str, field_schema: str
|
||||
) -> None: ...
|
||||
|
||||
def upsert(self, name: str, points: Sequence[VectorPoint]) -> int: ...
|
||||
|
||||
def count(self, name: str) -> int: ...
|
||||
|
||||
def retrieve(self, name: str, point_id: str) -> Optional[VectorPoint]: ...
|
||||
|
||||
def find_by_payload(
|
||||
self, name: str, equals: Mapping[str, Any]
|
||||
) -> Sequence[VectorPoint]:
|
||||
"""Every point whose payload matches all of `equals`. No vector, no ranking.
|
||||
|
||||
This is mode A of `docs/v1-delivery-plan.md` §3, and it is a `scroll`
|
||||
rather than a `search` on purpose. The rule the plan refuses to bend is
|
||||
"return the whole section, not a top-k of fragments" — returning two of
|
||||
five contraindications is more dangerous than returning none, because a
|
||||
partial list reads as a complete one. A ranked search cannot express
|
||||
that; an exhaustive filter can.
|
||||
|
||||
A list-valued field (`atc_codes`) matches when any element equals the
|
||||
given value.
|
||||
"""
|
||||
...
|
||||
@@ -0,0 +1,163 @@
|
||||
"""The only module in this package that knows `qdrant_client` exists.
|
||||
|
||||
Same arrangement as `embed/bedrock_runtime.py` and for the same reason: the
|
||||
loader's rules — corpus binding, derived point ids, batching, the point-count
|
||||
gate — are exercised in full against `InMemoryVectorStore` with no server
|
||||
running, and this file is the thin edge where those rules meet a real database.
|
||||
|
||||
The import sits inside `_client()` so `ingestion.load` imports cleanly on a
|
||||
machine with no `qdrant-client` installed and no Qdrant reachable, which is the
|
||||
state this repository is in by default.
|
||||
|
||||
What this adapter deliberately does not do is retry, shard, or tune. A load is
|
||||
an offline batch run a human starts and watches; a failure should surface, not
|
||||
be smoothed over into a partially loaded collection.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, List, Mapping, Optional, Sequence
|
||||
|
||||
from .models import CollectionSpec, VectorPoint
|
||||
|
||||
DEFAULT_URL = "http://localhost:6333"
|
||||
|
||||
# Points fetched per scroll round-trip. Only affects round trips, never the
|
||||
# result: `find_by_payload` pages until the collection says there is no more.
|
||||
SCROLL_PAGE = 256
|
||||
|
||||
_DISTANCES = {"cosine": "Cosine", "dot": "Dot", "euclid": "Euclid"}
|
||||
|
||||
|
||||
class QdrantVectorStore:
|
||||
def __init__(
|
||||
self,
|
||||
url: str = DEFAULT_URL,
|
||||
api_key: Optional[str] = None,
|
||||
timeout: float = 60.0,
|
||||
client: Optional[Any] = None,
|
||||
) -> None:
|
||||
self._url = url
|
||||
self._api_key = api_key
|
||||
self._timeout = timeout
|
||||
self._client = client
|
||||
|
||||
@property
|
||||
def url(self) -> str:
|
||||
return self._url
|
||||
|
||||
def _client_or_connect(self) -> Any:
|
||||
if self._client is None:
|
||||
from qdrant_client import QdrantClient
|
||||
|
||||
self._client = QdrantClient(
|
||||
url=self._url, api_key=self._api_key, timeout=self._timeout
|
||||
)
|
||||
return self._client
|
||||
|
||||
def collection_exists(self, name: str) -> bool:
|
||||
client = self._client_or_connect()
|
||||
existing = client.get_collections().collections
|
||||
return any(collection.name == name for collection in existing)
|
||||
|
||||
def create_collection(self, spec: CollectionSpec) -> None:
|
||||
from qdrant_client.models import Distance, VectorParams
|
||||
|
||||
distance = _DISTANCES.get(spec.distance.lower())
|
||||
if distance is None:
|
||||
raise ValueError(
|
||||
f"unsupported distance {spec.distance!r}; "
|
||||
f"expected one of {sorted(_DISTANCES)}"
|
||||
)
|
||||
self._client_or_connect().create_collection(
|
||||
collection_name=spec.name,
|
||||
vectors_config=VectorParams(
|
||||
size=spec.vector_size, distance=Distance(distance)
|
||||
),
|
||||
)
|
||||
|
||||
def delete_collection(self, name: str) -> None:
|
||||
self._client_or_connect().delete_collection(collection_name=name)
|
||||
|
||||
def create_payload_index(
|
||||
self, name: str, field_name: str, field_schema: str
|
||||
) -> None:
|
||||
self._client_or_connect().create_payload_index(
|
||||
collection_name=name,
|
||||
field_name=field_name,
|
||||
field_schema=field_schema,
|
||||
)
|
||||
|
||||
def upsert(self, name: str, points: Sequence[VectorPoint]) -> int:
|
||||
from qdrant_client.models import PointStruct
|
||||
|
||||
if not points:
|
||||
return 0
|
||||
self._client_or_connect().upsert(
|
||||
collection_name=name,
|
||||
points=[
|
||||
PointStruct(id=point.id, vector=point.vector, payload=point.payload)
|
||||
for point in points
|
||||
],
|
||||
wait=True,
|
||||
)
|
||||
return len(points)
|
||||
|
||||
def count(self, name: str) -> int:
|
||||
return self._client_or_connect().count(
|
||||
collection_name=name, exact=True
|
||||
).count
|
||||
|
||||
def find_by_payload(
|
||||
self, name: str, equals: Mapping[str, Any]
|
||||
) -> List[VectorPoint]:
|
||||
from qdrant_client.models import FieldCondition, Filter, MatchValue
|
||||
|
||||
if not equals:
|
||||
raise ValueError("find_by_payload needs at least one condition")
|
||||
scroll_filter = Filter(
|
||||
must=[
|
||||
FieldCondition(key=key, match=MatchValue(value=value))
|
||||
for key, value in equals.items()
|
||||
]
|
||||
)
|
||||
client = self._client_or_connect()
|
||||
found: List[VectorPoint] = []
|
||||
offset = None
|
||||
while True:
|
||||
page, offset = client.scroll(
|
||||
collection_name=name,
|
||||
scroll_filter=scroll_filter,
|
||||
limit=SCROLL_PAGE,
|
||||
offset=offset,
|
||||
with_payload=True,
|
||||
with_vectors=False,
|
||||
)
|
||||
found.extend(
|
||||
VectorPoint(
|
||||
id=str(record.id),
|
||||
vector=[],
|
||||
payload=dict(record.payload or {}),
|
||||
)
|
||||
for record in page
|
||||
)
|
||||
# Paging must not stop early. A section split into more parts than
|
||||
# one page holds would come back truncated, and a truncated section
|
||||
# is the one failure mode mode A exists to prevent.
|
||||
if offset is None:
|
||||
return found
|
||||
|
||||
def retrieve(self, name: str, point_id: str) -> Optional[VectorPoint]:
|
||||
found = self._client_or_connect().retrieve(
|
||||
collection_name=name,
|
||||
ids=[point_id],
|
||||
with_payload=True,
|
||||
with_vectors=True,
|
||||
)
|
||||
if not found:
|
||||
return None
|
||||
record = found[0]
|
||||
return VectorPoint(
|
||||
id=str(record.id),
|
||||
vector=list(record.vector or []),
|
||||
payload=dict(record.payload or {}),
|
||||
)
|
||||
@@ -0,0 +1,124 @@
|
||||
"""Corpus embed + load entry point.
|
||||
|
||||
`cli.py` belongs to the parser/chunking work, so this is the standalone entry
|
||||
point promised to Codex in `coordination/CLAUDE_TASK_2026-08-04.md` §4.1:
|
||||
|
||||
python -m ingestion.load.run --provider cohere-v4 --collection duocthu_v1
|
||||
|
||||
Two phases, deliberately separable. Embedding is the only part that leaves the
|
||||
machine, so it goes through the disk cache: an interrupted run resumes from
|
||||
whatever it already paid for instead of re-embedding it. Loading then reads
|
||||
that cache and never calls a provider at all.
|
||||
|
||||
Vectors are keyed by `(model_id, input_kind, sha256(text))`, so a re-run after
|
||||
an unrelated chunk edit re-embeds only the texts that actually changed.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Iterator, List, Sequence, Tuple
|
||||
|
||||
from ..embed.cache import CachingEmbeddingProvider, EmbeddingCache
|
||||
from ..embed.ports import INPUT_DOCUMENT
|
||||
from ..embed.registry import DEFAULT_REGION, build_provider
|
||||
from .corpus import corpus_sha256, count_chunks, iter_chunk_records
|
||||
from .models import CollectionSpec, CorpusManifest
|
||||
from .qdrant_repo import QdrantVectorStore
|
||||
from .upsert import ChunkLoader
|
||||
|
||||
DEFAULT_SLICE = 960
|
||||
DEFAULT_ATTEMPTS = 3
|
||||
|
||||
|
||||
def _slices(items: Sequence, size: int) -> Iterator[Tuple[int, Sequence]]:
|
||||
for start in range(0, len(items), size):
|
||||
yield start, items[start : start + size]
|
||||
|
||||
|
||||
def _embed_with_retry(provider, texts: Sequence[str], attempts: int) -> List:
|
||||
last: Exception | None = None
|
||||
for attempt in range(1, attempts + 1):
|
||||
try:
|
||||
return list(provider.embed(texts, INPUT_DOCUMENT).vectors)
|
||||
except Exception as exc: # noqa: BLE001 — a flaky link is the norm here
|
||||
last = exc
|
||||
if attempt == attempts:
|
||||
break
|
||||
backoff = 2**attempt
|
||||
print(f" attempt {attempt} failed ({exc}); retrying in {backoff}s")
|
||||
time.sleep(backoff)
|
||||
raise RuntimeError(f"embedding failed after {attempts} attempts") from last
|
||||
|
||||
|
||||
def main(argv=None) -> int:
|
||||
parser = argparse.ArgumentParser(prog="python -m ingestion.load.run")
|
||||
parser.add_argument("--chunks", type=Path, default=Path("data/processed/chunks.jsonl"))
|
||||
parser.add_argument("--cache", type=Path, default=Path("data/processed/embeddings"))
|
||||
parser.add_argument("--provider", required=True)
|
||||
parser.add_argument("--collection", required=True)
|
||||
parser.add_argument("--region", default=DEFAULT_REGION)
|
||||
parser.add_argument("--qdrant-url", default="http://localhost:6333")
|
||||
parser.add_argument("--slice-size", type=int, default=DEFAULT_SLICE)
|
||||
parser.add_argument("--attempts", type=int, default=DEFAULT_ATTEMPTS)
|
||||
parser.add_argument("--embed-only", action="store_true")
|
||||
args = parser.parse_args(argv)
|
||||
|
||||
records = list(iter_chunk_records(args.chunks))
|
||||
sha = corpus_sha256(args.chunks)
|
||||
print(f"corpus : {args.chunks}")
|
||||
print(f"chunks : {len(records)} (count_chunks={count_chunks(args.chunks)})")
|
||||
print(f"sha256 : {sha}")
|
||||
|
||||
inner = build_provider(args.provider, region=args.region)
|
||||
cache_path = args.cache / f"{args.provider}.jsonl"
|
||||
cache_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
provider = CachingEmbeddingProvider(inner, EmbeddingCache(cache_path))
|
||||
print(f"provider : {provider.model_id} ({provider.dimensions}d)")
|
||||
print(f"cache : {cache_path}")
|
||||
|
||||
texts = [str(record["text"]) for record in records]
|
||||
vectors: List = []
|
||||
started = time.time()
|
||||
for start, slice_texts in _slices(texts, args.slice_size):
|
||||
vectors.extend(_embed_with_retry(provider, slice_texts, args.attempts))
|
||||
done = len(vectors)
|
||||
rate = done / max(time.time() - started, 1e-9)
|
||||
remaining = (len(texts) - done) / rate if rate else 0
|
||||
print(
|
||||
f" embedded {done}/{len(texts)} "
|
||||
f"({rate:.1f}/s, ~{remaining/60:.1f} min left)",
|
||||
flush=True,
|
||||
)
|
||||
|
||||
stats = provider.cache.stats
|
||||
print(f"cache : {stats.hits} hits, {stats.misses} misses")
|
||||
|
||||
if args.embed_only:
|
||||
print("embed-only: stopping before the vector store")
|
||||
return 0
|
||||
|
||||
manifest = CorpusManifest(
|
||||
corpus_sha256=sha,
|
||||
chunk_count=len(records),
|
||||
model_id=provider.model_id,
|
||||
dimensions=provider.dimensions,
|
||||
input_kind=INPUT_DOCUMENT,
|
||||
provider=args.provider,
|
||||
)
|
||||
spec = CollectionSpec(name=args.collection, vector_size=provider.dimensions)
|
||||
store = QdrantVectorStore(url=args.qdrant_url)
|
||||
loader = ChunkLoader(store, spec, manifest)
|
||||
|
||||
report = loader.load(zip(records, (v.values for v in vectors)))
|
||||
print(
|
||||
f"loaded : {report.points_upserted} points in {report.batches} batches; "
|
||||
f"collection holds {report.collection_count}; "
|
||||
f"count gate {'PASS' if report.count_matches else 'FAIL'}"
|
||||
)
|
||||
return 0 if report.count_matches else 1
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@@ -0,0 +1,129 @@
|
||||
"""Loading chunks into a collection, idempotently and only into the right one.
|
||||
|
||||
The loader does three things in a fixed order, and the order is the point.
|
||||
It checks the corpus binding *before* creating or writing anything (A6), so a
|
||||
refused load leaves the store untouched rather than half-overwritten. It derives
|
||||
every point id from `chunk_id` (A5), so the same corpus loaded twice converges
|
||||
instead of doubling. And it reports the collection's point count against the
|
||||
number of chunks it was given, which is the v1 gate
|
||||
`qdrant_point_count != chunk_count`.
|
||||
|
||||
Vectors are validated against the collection's declared size before the first
|
||||
upsert. A wrong-sized vector is a whole-run defect, not a bad record, and
|
||||
failing on the first one is cheaper than discovering it after 15,066 upserts.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Iterable, List, Mapping, Sequence, Tuple
|
||||
|
||||
from .manifest import assert_compatible, write_manifest
|
||||
from .models import CollectionSpec, CorpusManifest, VectorPoint, build_point
|
||||
from .ports import VectorStore
|
||||
|
||||
ChunkVectorPair = Tuple[Mapping[str, Any], Sequence[float]]
|
||||
|
||||
DEFAULT_BATCH_SIZE = 256
|
||||
|
||||
|
||||
class PointCountMismatch(RuntimeError):
|
||||
"""The collection does not hold exactly one point per chunk."""
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class LoadReport:
|
||||
collection: str
|
||||
collection_created: bool
|
||||
points_upserted: int
|
||||
batches: int
|
||||
collection_count: int
|
||||
corpus_sha256: str
|
||||
|
||||
@property
|
||||
def count_matches(self) -> bool:
|
||||
return self.collection_count == self.points_upserted
|
||||
|
||||
|
||||
class ChunkLoader:
|
||||
def __init__(
|
||||
self,
|
||||
store: VectorStore,
|
||||
spec: CollectionSpec,
|
||||
manifest: CorpusManifest,
|
||||
batch_size: int = DEFAULT_BATCH_SIZE,
|
||||
) -> None:
|
||||
if batch_size <= 0:
|
||||
raise ValueError(f"batch_size must be positive, got {batch_size}")
|
||||
if manifest.dimensions != spec.vector_size:
|
||||
raise ValueError(
|
||||
f"manifest declares {manifest.dimensions} dimensions but the "
|
||||
f"collection spec declares {spec.vector_size}"
|
||||
)
|
||||
self._store = store
|
||||
self._spec = spec
|
||||
self._manifest = manifest
|
||||
self._batch_size = batch_size
|
||||
|
||||
def prepare(self) -> bool:
|
||||
"""Gate on the corpus binding, then ensure the collection exists.
|
||||
|
||||
Returns True when the collection was created by this call.
|
||||
"""
|
||||
assert_compatible(self._store, self._spec.name, self._manifest)
|
||||
|
||||
created = False
|
||||
if not self._store.collection_exists(self._spec.name):
|
||||
self._store.create_collection(self._spec)
|
||||
for field_name, field_schema in self._spec.indexed_fields:
|
||||
self._store.create_payload_index(
|
||||
self._spec.name, field_name, field_schema
|
||||
)
|
||||
created = True
|
||||
write_manifest(self._store, self._spec.name, self._manifest)
|
||||
return created
|
||||
|
||||
def load(self, pairs: Iterable[ChunkVectorPair]) -> LoadReport:
|
||||
created = self.prepare()
|
||||
|
||||
upserted = 0
|
||||
batches = 0
|
||||
buffer: List[VectorPoint] = []
|
||||
for record, vector in pairs:
|
||||
buffer.append(self._point(record, vector))
|
||||
if len(buffer) >= self._batch_size:
|
||||
upserted += self._store.upsert(self._spec.name, buffer)
|
||||
batches += 1
|
||||
buffer = []
|
||||
if buffer:
|
||||
upserted += self._store.upsert(self._spec.name, buffer)
|
||||
batches += 1
|
||||
|
||||
return LoadReport(
|
||||
collection=self._spec.name,
|
||||
collection_created=created,
|
||||
points_upserted=upserted,
|
||||
batches=batches,
|
||||
collection_count=self._store.count(self._spec.name),
|
||||
corpus_sha256=self._manifest.corpus_sha256,
|
||||
)
|
||||
|
||||
def assert_point_count(self, expected_chunks: int) -> int:
|
||||
"""v1 gate: exactly one point per chunk, no more and no fewer."""
|
||||
actual = self._store.count(self._spec.name)
|
||||
if actual != expected_chunks:
|
||||
raise PointCountMismatch(
|
||||
f"collection {self._spec.name!r} holds {actual} points but the "
|
||||
f"corpus has {expected_chunks} chunks"
|
||||
)
|
||||
return actual
|
||||
|
||||
def _point(
|
||||
self, record: Mapping[str, Any], vector: Sequence[float]
|
||||
) -> VectorPoint:
|
||||
if len(vector) != self._spec.vector_size:
|
||||
raise ValueError(
|
||||
f"chunk {record.get('chunk_id', '<no chunk_id>')!r} has a "
|
||||
f"{len(vector)}-dimension vector, collection "
|
||||
f"{self._spec.name!r} expects {self._spec.vector_size}"
|
||||
)
|
||||
return build_point(record, vector)
|
||||
@@ -30,7 +30,7 @@ from typing import Iterator, List, Optional, Union
|
||||
from ..extract.models import Span
|
||||
from ..extract.page_map import HEADER_BAND_Y
|
||||
from ..normalize import join_spans, substitute_pua
|
||||
from ..tables.classify import QUARANTINE_SHAPES
|
||||
from ..tables.classify import QUARANTINE_SHAPES, SHAPE_FORMULA_2D
|
||||
from .atc import extract_atc_codes
|
||||
from .detector import in_monograph_range, is_monograph_title_candidate
|
||||
from .merge import merge_multiline_headings, merge_same_line_bold_fragments
|
||||
@@ -51,6 +51,9 @@ from .vocab import (
|
||||
)
|
||||
|
||||
_QUALIFIER_RE = re.compile(r"^\(.+\)$")
|
||||
_DOSING_TABLE_CAPTION_RE = re.compile(
|
||||
r"(?:^|\b)Bảng\s+\d+\s*[.:]\s*Điều chỉnh liều\b", re.IGNORECASE,
|
||||
)
|
||||
|
||||
|
||||
def _is_page_boilerplate(span: Span) -> bool:
|
||||
@@ -94,6 +97,31 @@ def _slugify(text: str) -> str:
|
||||
return re.sub(r"[^a-z0-9]+", "_", ascii_text.lower()).strip("_")
|
||||
|
||||
|
||||
def _starts_its_visual_line(span: Span, previous: Span | None) -> bool:
|
||||
"""True when nothing else was printed to the left of this span on its line.
|
||||
|
||||
A real section heading opens a line. Confirmed content loss when this was
|
||||
not checked: CISPLATIN (physical page 402) prints `Suy thận: Chống chỉ
|
||||
định.` inside its dosing section, and the second half is a section name.
|
||||
Matched as a heading, `Chống chỉ định.` vanished from the dosing text and
|
||||
the section ended on a bare `Suy thận:` — a renal-impairment
|
||||
contraindication silently dropped. ISOPRENALIN had the same shape.
|
||||
|
||||
"To the left" is decided on `x0`, not on `previous.x1 <= span.x0`: NEVIRAPIN
|
||||
(physical page 1045) prints `Xem thêm mục ` at x1=104.89 immediately before
|
||||
an italic `Liều lượng và cách dùng` at x0=104.88 — a 0.01pt overlap from the
|
||||
trailing space's advance width, enough to make an end-before-start test call
|
||||
a mid-line cross-reference a heading. Comparing the left edges cannot be
|
||||
defeated by glyph-advance rounding and still reads False for the synthetic
|
||||
fixtures, which place every span at identical coordinates.
|
||||
"""
|
||||
if previous is None:
|
||||
return True
|
||||
same_line = (previous.physical_page, previous.block, previous.line) == (
|
||||
span.physical_page, span.block, span.line)
|
||||
return not (same_line and previous.x0 < span.x0)
|
||||
|
||||
|
||||
def _is_body_line_that_reads_like_a_label(span: Span, items: List) -> bool:
|
||||
"""A plain line that repeats a section name, sitting under a heading.
|
||||
|
||||
@@ -114,7 +142,48 @@ def _is_body_line_that_reads_like_a_label(span: Span, items: List) -> bool:
|
||||
return bool(items) and isinstance(items[-1], _SectionEvent)
|
||||
|
||||
|
||||
def _classify(spans: List[Span]) -> List[Union[Span, _SectionEvent, _TextEvent]]:
|
||||
def _is_mid_line_label(span: Span, previous: Span | None) -> bool:
|
||||
"""A non-bold section name printed part-way along a line is body text."""
|
||||
return not span.bold and not _starts_its_visual_line(span, previous)
|
||||
|
||||
|
||||
def _continues_previous_visual_line(span: Span, previous: Span | None) -> bool:
|
||||
"""A plain label-shaped span can be the wrapped tail of body prose.
|
||||
|
||||
Confirmed across the corpus: sentences such as ``không phải là`` /
|
||||
``chống chỉ định.`` remain in the same PDF block on adjacent visual
|
||||
lines. Exact vocabulary matching used to consume the second line as a
|
||||
heading. A real plain heading observed in this book follows completed
|
||||
prose; a non-bold adjacent continuation after an unterminated line does
|
||||
not. Styling and PDF block geometry make this deliberately narrower than
|
||||
a text-only language heuristic.
|
||||
"""
|
||||
if span.bold or previous is None:
|
||||
return False
|
||||
same_block = (previous.physical_page, previous.block) == (
|
||||
span.physical_page, span.block)
|
||||
adjacent_line = span.line == previous.line + 1
|
||||
previous_text = previous.text.rstrip()
|
||||
terminal = previous_text.endswith((".", "!", "?", ":", ";"))
|
||||
return same_block and adjacent_line and bool(previous_text) and not terminal
|
||||
|
||||
|
||||
def _is_italic_cross_reference(span: Span) -> bool:
|
||||
"""The book italicises references to other sections; headings are never italic.
|
||||
|
||||
Measured over the whole monograph range (physical 99-1496): 11,916 spans
|
||||
carrying a section name are bold, 29 are plain, and **7 are italic — none
|
||||
of them a heading**. Two of the seven wrap onto a line of their own, where
|
||||
the "opens its line" test cannot help: CALCI LACTAT (p296) breaks `xem thêm
|
||||
về nhu cầu hàng ngày trong mục ` / `Dược lý và cơ chế tác dụng)`, and
|
||||
physical page 432 breaks ordinary prose — `hoặc khi không được ` /
|
||||
`chỉ định.` — across a line, so `chỉ định.` alone opened a spurious
|
||||
"Chỉ định" section.
|
||||
"""
|
||||
return span.italic
|
||||
|
||||
|
||||
def _classify(spans: List[Span], table_index=None) -> List[Union[Span, _SectionEvent, _TextEvent]]:
|
||||
"""Pass 1: tag each span. Title candidates are left as raw Span objects
|
||||
(pass 2 groups + merges them); everything else becomes a typed event.
|
||||
|
||||
@@ -127,29 +196,110 @@ def _classify(spans: List[Span]) -> List[Union[Span, _SectionEvent, _TextEvent]]
|
||||
same lesson as "don't gate on font size" (ADR 0003 item 10) applied to
|
||||
boldness instead.
|
||||
"""
|
||||
# PyMuPDF's block order is not guaranteed to keep every cell of a table
|
||||
# together. A visually later section heading can therefore occur between
|
||||
# cells of one physical region in the extracted stream (confirmed on
|
||||
# CAPECITABIN pp. 308-309 and IMATINIB p. 795). Gather each region first,
|
||||
# then emit it atomically at its first occurrence. Besides preserving the
|
||||
# section active at the top of the table, this guarantees one run/block per
|
||||
# physical region instead of duplicate fragments with conflicting owners.
|
||||
region_spans: dict[int, List[Span]] = {}
|
||||
region_for_span: dict[int, object] = {}
|
||||
spans_by_page: dict[int, List[Span]] = {}
|
||||
for span in spans:
|
||||
spans_by_page.setdefault(span.physical_page, []).append(span)
|
||||
region = _region_for(table_index, span)
|
||||
if region is None:
|
||||
continue
|
||||
region_for_span[id(span)] = region
|
||||
region_spans.setdefault(id(region), []).append(span)
|
||||
|
||||
# A table can continue after the monograph's ordinary final headings.
|
||||
# CAPECITABIN p. 309 does exactly that: two dose-adjustment tables follow
|
||||
# the trade-name line, with their own explicit captions but without a
|
||||
# repeated "Liều lượng và cách dùng" heading. Use only the narrow,
|
||||
# unambiguous caption signal and its geometry; a generic keyword rule
|
||||
# would be unsafe in clinical prose.
|
||||
caption_for_region: dict[int, Span] = {}
|
||||
region_for_caption: dict[int, object] = {}
|
||||
for page_regions in (table_index or {}).values():
|
||||
for region in page_regions:
|
||||
candidates = [
|
||||
span for span in spans_by_page.get(region.physical_page, ())
|
||||
if 0 <= region.bbox[1] - span.y1 <= 80
|
||||
and _DOSING_TABLE_CAPTION_RE.search(span.text.strip())
|
||||
]
|
||||
if candidates:
|
||||
caption = max(candidates, key=lambda candidate: candidate.y1)
|
||||
caption_for_region[id(region)] = caption
|
||||
region_for_caption[id(caption)] = region
|
||||
|
||||
items: List[Union[Span, _SectionEvent, _TextEvent]] = []
|
||||
emitted_regions: set[int] = set()
|
||||
emitted_captions: set[int] = set()
|
||||
previous: Span | None = None
|
||||
for span in spans:
|
||||
if not span.text.strip():
|
||||
continue
|
||||
if _is_page_boilerplate(span):
|
||||
previous = span
|
||||
continue
|
||||
caption_region = region_for_caption.get(id(span))
|
||||
if caption_region is not None:
|
||||
caption_key = id(span)
|
||||
if caption_key not in emitted_captions:
|
||||
dose_section = match_section("Liều lượng và cách dùng")
|
||||
assert dose_section is not None
|
||||
items.append(_SectionEvent(dose_section, span, inline_value=span.text))
|
||||
emitted_captions.add(caption_key)
|
||||
previous = span
|
||||
continue
|
||||
# Out-of-scope index text must never create section events, and cells
|
||||
# inside a known table are content even when a cell says "Chỉ định".
|
||||
# Classification previously happened without either context, letting
|
||||
# the WARFARIN/IOBITRIDOL table header change the owning section.
|
||||
region = region_for_span.get(id(span))
|
||||
if region is not None:
|
||||
region_key = id(region)
|
||||
if region_key not in emitted_regions:
|
||||
caption = caption_for_region.get(region_key)
|
||||
if caption is not None and id(caption) not in emitted_captions:
|
||||
dose_section = match_section("Liều lượng và cách dùng")
|
||||
assert dose_section is not None
|
||||
items.append(_SectionEvent(dose_section, caption,
|
||||
inline_value=caption.text))
|
||||
emitted_captions.add(id(caption))
|
||||
items.extend(_TextEvent(cell) for cell in region_spans[region_key]
|
||||
if cell.text.strip() and not _is_page_boilerplate(cell))
|
||||
emitted_regions.add(region_key)
|
||||
previous = span
|
||||
continue
|
||||
if not in_monograph_range(span):
|
||||
items.append(_TextEvent(span))
|
||||
previous = span
|
||||
continue
|
||||
if is_monograph_title_candidate(span):
|
||||
items.append(span)
|
||||
previous = span
|
||||
continue
|
||||
section_def = match_section(span.text)
|
||||
if section_def is not None and not _is_body_line_that_reads_like_a_label(
|
||||
span, items
|
||||
):
|
||||
items.append(_SectionEvent(section_def, span))
|
||||
continue
|
||||
if section_def is not None:
|
||||
items.append(_TextEvent(span))
|
||||
is_body = (_is_body_line_that_reads_like_a_label(span, items)
|
||||
or _is_mid_line_label(span, previous)
|
||||
or _continues_previous_visual_line(span, previous)
|
||||
or _is_italic_cross_reference(span))
|
||||
items.append(_TextEvent(span) if is_body
|
||||
else _SectionEvent(section_def, span))
|
||||
previous = span
|
||||
continue
|
||||
inline = match_section_with_inline_value(span.text)
|
||||
if inline is not None:
|
||||
if (inline is not None and not _is_mid_line_label(span, previous)
|
||||
and not _continues_previous_visual_line(span, previous)
|
||||
and not _is_italic_cross_reference(span)):
|
||||
items.append(_SectionEvent(inline[0], span, inline_value=inline[1]))
|
||||
else:
|
||||
items.append(_TextEvent(span))
|
||||
previous = span
|
||||
return items
|
||||
|
||||
|
||||
@@ -228,6 +378,17 @@ def _region_for(table_index, span: Span):
|
||||
if not table_index:
|
||||
return None
|
||||
for region in table_index.get(span.physical_page, ()):
|
||||
if region.shape == SHAPE_FORMULA_2D and span.column in ("left", "right"):
|
||||
# Formula bands are intentionally widened enough to reach past
|
||||
# the gutter (some numerator spans have misleading leading-space
|
||||
# boxes), so geometry alone can swallow prose from the opposite
|
||||
# column. Keep the wide band but require its source column to
|
||||
# agree with the span's extracted block column. This preserves
|
||||
# AMPICILIN's gutter-adjacent "Cl" while excluding NETILMICIN's
|
||||
# right-column cross-reference from a left-column formula.
|
||||
region_column = "left" if (region.bbox[0] + region.bbox[2]) / 2 < 303.5 else "right"
|
||||
if span.column != region_column:
|
||||
continue
|
||||
if region.contains(span.x0, span.y0, span.x1, span.y1):
|
||||
return region
|
||||
return None
|
||||
@@ -259,7 +420,7 @@ def assemble(spans: List[Span], table_index=None, ledger: Optional[list] = None)
|
||||
"""
|
||||
raw_chars = sum(len(s.text) for s in spans)
|
||||
spans = merge_same_line_bold_fragments(spans)
|
||||
events = _filter_false_positive_titles(_coalesce_titles(_classify(spans)))
|
||||
events = _filter_false_positive_titles(_coalesce_titles(_classify(spans, table_index)))
|
||||
|
||||
# Span-level coverage ledger. Character counts alone cannot balance here
|
||||
# (normalization joins, substitutes and drops characters), so every span
|
||||
@@ -291,7 +452,7 @@ def assemble(spans: List[Span], table_index=None, ledger: Optional[list] = None)
|
||||
current: Optional[Monograph] = None
|
||||
current_section_key: Optional[str] = None
|
||||
runs: List[tuple] = [] # ordered [(region_or_None, [spans])]
|
||||
inline_prefix: str = ""
|
||||
inline_prefix: Optional[tuple[str, Span]] = None
|
||||
awaiting_qualifier = False
|
||||
|
||||
def append_span(span: Span, region):
|
||||
@@ -338,10 +499,12 @@ def assemble(spans: List[Span], table_index=None, ledger: Optional[list] = None)
|
||||
if region is not None else False,
|
||||
))
|
||||
if inline_prefix:
|
||||
inline_text, inline_span = inline_prefix
|
||||
head = SectionPart(
|
||||
kind=PART_PROSE, text=inline_prefix,
|
||||
physical_page=parts[0].physical_page if parts else 0,
|
||||
bbox=parts[0].bbox if parts else [0.0, 0.0, 0.0, 0.0],
|
||||
kind=PART_PROSE, text=inline_text,
|
||||
physical_page=inline_span.physical_page,
|
||||
bbox=[inline_span.x0, inline_span.y0, inline_span.x1, inline_span.y1],
|
||||
source_span_ids=[inline_span.span_id],
|
||||
)
|
||||
parts.insert(0, head)
|
||||
return parts
|
||||
@@ -387,7 +550,7 @@ def assemble(spans: List[Span], table_index=None, ledger: Optional[list] = None)
|
||||
source_span_ids=list(part.source_span_ids),
|
||||
))
|
||||
runs = []
|
||||
inline_prefix = ""
|
||||
inline_prefix = None
|
||||
|
||||
def finalize(monograph: Monograph) -> Monograph:
|
||||
# Duplicate check happens here, not at title-detection time: the
|
||||
@@ -401,8 +564,14 @@ def assemble(spans: List[Span], table_index=None, ledger: Optional[list] = None)
|
||||
f"line (outlier item 18) before assuming this is a real collision"
|
||||
)
|
||||
seen_ids.add(monograph.drug_id)
|
||||
if "ma_atc" in monograph.sections:
|
||||
result = extract_atc_codes(monograph.sections["ma_atc"].text)
|
||||
atc_source = monograph.sections.get("ma_atc")
|
||||
# One real class monograph combines the two labels as "Tên chung
|
||||
# quốc tế và mã ATC". It remains the required title anchor, and
|
||||
# its ATC codes are still extracted rather than silently discarded.
|
||||
if atc_source is None:
|
||||
atc_source = monograph.sections.get("ten_chung_quoc_te")
|
||||
if atc_source is not None:
|
||||
result = extract_atc_codes(atc_source.text)
|
||||
monograph.atc_codes = result.codes
|
||||
monograph.atc_stated_absent = result.stated_absent
|
||||
return monograph
|
||||
@@ -440,7 +609,7 @@ def assemble(spans: List[Span], table_index=None, ledger: Optional[list] = None)
|
||||
text="",
|
||||
)
|
||||
if event.inline_value:
|
||||
inline_prefix = event.inline_value
|
||||
inline_prefix = (event.inline_value, event.span)
|
||||
awaiting_qualifier = False
|
||||
mark(event.span, SPAN_STATE_HEADING)
|
||||
current.source_page_range[1] = max(current.source_page_range[1], event.span.physical_page)
|
||||
|
||||
@@ -30,6 +30,13 @@ from .vocab import is_part_divider, match_section
|
||||
|
||||
MONOGRAPH_PRINTED_PAGE_START = 99
|
||||
MONOGRAPH_PRINTED_PAGE_END = 1496
|
||||
# Printed folios are metadata inferred from page headers and can be wrong in
|
||||
# the back index (confirmed: physical page 1655 was inferred as printed page
|
||||
# 1496). The physical bounds are source-document invariants and prevent an
|
||||
# index entry that looks like a section label from extending ZOLPIDEM by 161
|
||||
# pages. Keep both checks: either signal alone has known failure modes.
|
||||
MONOGRAPH_PHYSICAL_PAGE_START = 99
|
||||
MONOGRAPH_PHYSICAL_PAGE_END = 1496
|
||||
_MIN_TITLE_LEN = 3
|
||||
_MAX_TITLE_LEN = 60
|
||||
_MAX_LOWERCASE_RATIO = 0.10 # HMG-CoA: 1/27 = 3.7% (real title) vs "Mã ATC:": 1/5 = 20% (real
|
||||
@@ -49,6 +56,8 @@ def _is_mostly_upper(text: str) -> bool:
|
||||
|
||||
def in_monograph_range(span: Span) -> bool:
|
||||
return (
|
||||
MONOGRAPH_PHYSICAL_PAGE_START <= span.physical_page <= MONOGRAPH_PHYSICAL_PAGE_END
|
||||
and
|
||||
span.printed_page is not None
|
||||
and MONOGRAPH_PRINTED_PAGE_START <= span.printed_page <= MONOGRAPH_PRINTED_PAGE_END
|
||||
)
|
||||
|
||||
@@ -37,11 +37,12 @@ class SectionDef:
|
||||
|
||||
|
||||
SECTION_DEFS = [
|
||||
SectionDef("ten_chung_quoc_te", "Tên chung quốc tế", ("Ten chung quốc tế",)),
|
||||
SectionDef("ten_chung_quoc_te", "Tên chung quốc tế",
|
||||
("Ten chung quốc tế", "Tên chung quốc tế và mã ATC")),
|
||||
SectionDef("ma_atc", "Mã ATC", ("Mã ACT",)),
|
||||
SectionDef("loai_thuoc", "Loại thuốc", ("Loại thuôc", "Lọai thuốc", "Phân loại thuốc")),
|
||||
SectionDef("dang_thuoc_va_ham_luong", "Dạng thuốc và hàm lượng",
|
||||
("Dạng dùng và hàm lượng",)),
|
||||
("Dạng dùng và hàm lượng", "Dạng bào chế và hàm lượng")),
|
||||
SectionDef("duoc_ly_va_co_che_tac_dung", "Dược lý và cơ chế tác dụng",
|
||||
("Dược lí và cơ chế tác dụng", "Dược lý học và cơ chế tác dụng")),
|
||||
SectionDef("chi_dinh", "Chỉ định"),
|
||||
@@ -56,7 +57,8 @@ SECTION_DEFS = [
|
||||
"Hướng dẫn cách xử trí các ADR")),
|
||||
SectionDef("lieu_luong_va_cach_dung", "Liều lượng và cách dùng",
|
||||
("Liều lượng cách dùng", "Liều lượng, cách dùng",
|
||||
"Liều dùng và cách dùng", "Liều lượng và cách sử dụng")),
|
||||
"Liều dùng và cách dùng", "Liều lượng và cách sử dụng",
|
||||
"Liều lượng và cách dùng giải độc tố uốn ván hấp phụ đơn giá")),
|
||||
SectionDef("tuong_tac_thuoc", "Tương tác thuốc"),
|
||||
SectionDef("do_on_dinh_va_bao_quan", "Độ ổn định và bảo quản"),
|
||||
SectionDef("tuong_ky", "Tương kỵ"),
|
||||
|
||||
@@ -8,6 +8,7 @@ from .readiness import (
|
||||
read_chunks,
|
||||
read_monographs,
|
||||
)
|
||||
from .clinical_readiness import evaluate_clinical
|
||||
from .residual_ink import (
|
||||
ANTIALIAS_SPECK,
|
||||
FRACTION_BAR_CANDIDATE,
|
||||
@@ -35,6 +36,7 @@ __all__ = [
|
||||
"read_chunks",
|
||||
"corpus_size",
|
||||
"read_monographs",
|
||||
"evaluate_clinical",
|
||||
"PageContext",
|
||||
"ResidualRegion",
|
||||
"classify",
|
||||
|
||||
@@ -23,7 +23,6 @@ from __future__ import annotations
|
||||
|
||||
import re
|
||||
from dataclasses import dataclass
|
||||
from typing import List
|
||||
|
||||
import fitz
|
||||
|
||||
@@ -32,6 +31,26 @@ BACK_INDEX_START_PHYSICAL = 1530 # printed 1531 — first page of real entries
|
||||
_ENTRY_RE = re.compile(r"^(.+?),\s*(\d+)\s*$")
|
||||
_CROSS_REF_MARKER = " - "
|
||||
|
||||
# Page furniture inside the index: the running header, the index's own title,
|
||||
# and the single-letter section dividers. These used to be dropped implicitly
|
||||
# by not matching the entry pattern — once wrapped lines are rejoined they
|
||||
# would instead be glued onto the entry below them, so they now have to be
|
||||
# named.
|
||||
_FURNITURE_RE = re.compile(r"^(DTQGVN\s*\d*|Mục lục tra cứu|[A-ZĐÀ-Ỹ])$")
|
||||
|
||||
# A bare number is ambiguous: the page's own folio, or the tail of an entry
|
||||
# whose page number wrapped onto the next line (`Bromhexine hydrochloride -
|
||||
# Bromhexin hydroclorid,` / `269`). Discarding it unconditionally cost real
|
||||
# entries — the unterminated fragment then swallowed the *following* entry, so
|
||||
# `Bromocriptin, 270` was consumed into a cross-reference and vanished from the
|
||||
# ground truth. It is furniture only when no fragment is waiting for it.
|
||||
_BARE_NUMBER_RE = re.compile(r"^\d{1,4}$")
|
||||
|
||||
# An index entry never wraps into a paragraph; the longest genuine one measured
|
||||
# in this book is well under this. A buffer that grows past it means the join
|
||||
# has lost the thread, so it is dropped rather than emitted as a bogus name.
|
||||
_MAX_JOINED_ENTRY_CHARS = 200
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class GroundTruthEntry:
|
||||
@@ -39,14 +58,82 @@ class GroundTruthEntry:
|
||||
printed_page: int
|
||||
|
||||
|
||||
def parse_back_index(doc: fitz.Document, start_physical_page: int = BACK_INDEX_START_PHYSICAL) -> List[GroundTruthEntry]:
|
||||
entries: List[GroundTruthEntry] = []
|
||||
@dataclass(frozen=True)
|
||||
class BackIndexAlias:
|
||||
alias: str
|
||||
target: str
|
||||
printed_page: int
|
||||
physical_page: int
|
||||
|
||||
|
||||
def _index_page_lines(doc: fitz.Document, start_physical_page: int):
|
||||
"""Index lines page by page, furniture removed. An entry never wraps across
|
||||
a page, so the caller can drop a pending fragment at each page break."""
|
||||
for pno in range(start_physical_page, doc.page_count):
|
||||
for line in doc[pno].get_text().split("\n"):
|
||||
line = line.strip()
|
||||
if not line or _CROSS_REF_MARKER in line:
|
||||
lines = [line.strip() for line in doc[pno].get_text().split("\n")]
|
||||
yield [line for line in lines if line and not _FURNITURE_RE.match(line)]
|
||||
|
||||
|
||||
def parse_back_index(doc: fitz.Document, start_physical_page: int = BACK_INDEX_START_PHYSICAL) -> list[GroundTruthEntry]:
|
||||
"""Entries in reading order, cross-references excluded.
|
||||
|
||||
Long entries wrap across two printed lines, and each fragment was read as
|
||||
an entry of its own before this was handled: `Acinet 10 - xem Atorvastatin
|
||||
- Các chất ức chế HMG - ` / `CoA reductase, 285` produced a phantom
|
||||
ground-truth drug called "- CoA reductase". Measured on the whole index,
|
||||
that shape accounted for 50 of the 76 entries `cli validate` could not
|
||||
match — noise that inflates the denominator and hides real misses inside
|
||||
it. A fragment is a line that does not yet end in ", <page>", so lines are
|
||||
accumulated until they do.
|
||||
"""
|
||||
entries: list[GroundTruthEntry] = []
|
||||
for page_lines in _index_page_lines(doc, start_physical_page):
|
||||
buffer = ""
|
||||
for line in page_lines:
|
||||
if not buffer and _BARE_NUMBER_RE.match(line):
|
||||
continue
|
||||
match = _ENTRY_RE.match(line)
|
||||
buffer = f"{buffer} {line}".strip() if buffer else line
|
||||
match = _ENTRY_RE.match(buffer)
|
||||
if match:
|
||||
entries.append(GroundTruthEntry(name=match.group(1).strip(), printed_page=int(match.group(2))))
|
||||
if _CROSS_REF_MARKER not in buffer:
|
||||
entries.append(GroundTruthEntry(
|
||||
name=match.group(1).strip(), printed_page=int(match.group(2))))
|
||||
buffer = ""
|
||||
elif len(buffer) > _MAX_JOINED_ENTRY_CHARS:
|
||||
buffer = ""
|
||||
return entries
|
||||
|
||||
|
||||
def parse_back_index_see_aliases(
|
||||
doc: fitz.Document,
|
||||
start_physical_page: int = BACK_INDEX_START_PHYSICAL,
|
||||
) -> list[BackIndexAlias]:
|
||||
"""Extract the book's explicit ``X - xem Y`` alias relations.
|
||||
|
||||
This is intentionally separate from :func:`parse_back_index`, whose
|
||||
validation semantics must remain unchanged.
|
||||
"""
|
||||
aliases: list[BackIndexAlias] = []
|
||||
marker = re.compile(r"\s+-\s+xem\s+", re.IGNORECASE)
|
||||
for physical_page, page_lines in enumerate(
|
||||
_index_page_lines(doc, start_physical_page), start_physical_page,
|
||||
):
|
||||
buffer = ""
|
||||
for line in page_lines:
|
||||
if not buffer and _BARE_NUMBER_RE.match(line):
|
||||
continue
|
||||
buffer = f"{buffer} {line}".strip() if buffer else line
|
||||
match = _ENTRY_RE.match(buffer)
|
||||
if match:
|
||||
parts = marker.split(match.group(1), maxsplit=1)
|
||||
if len(parts) == 2 and all(part.strip() for part in parts):
|
||||
aliases.append(BackIndexAlias(
|
||||
alias=parts[0].strip(),
|
||||
target=parts[1].strip(),
|
||||
printed_page=int(match.group(2)),
|
||||
physical_page=physical_page,
|
||||
))
|
||||
buffer = ""
|
||||
elif len(buffer) > _MAX_JOINED_ENTRY_CHARS:
|
||||
buffer = ""
|
||||
return aliases
|
||||
|
||||
@@ -0,0 +1,106 @@
|
||||
"""Executable gates for a clinical-production release.
|
||||
|
||||
These gates deliberately sit above parser/chunk readiness. Passing extraction
|
||||
tests cannot establish that the source is current, licensed, clinically
|
||||
reviewed, or safe to operate in a care setting.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
from pathlib import Path
|
||||
from typing import List
|
||||
|
||||
from .readiness import Gate
|
||||
|
||||
CURRENT_FORMULARY_EDITION = 3
|
||||
REQUIRED_APPROVAL_ROLES = {
|
||||
"physician",
|
||||
"clinical_pharmacist",
|
||||
"clinical_safety_owner",
|
||||
"regulatory_owner",
|
||||
}
|
||||
REQUIRED_RELEASE_ARTIFACTS = {
|
||||
"logical_tables": "data/processed/logical_tables.jsonl",
|
||||
"clinical_eval": "data/clinical/clinical_eval_report.json",
|
||||
"risk_management": "data/clinical/risk_management.json",
|
||||
"security_privacy": "data/clinical/security_privacy_review.json",
|
||||
"operations": "data/clinical/operations_readiness.json",
|
||||
}
|
||||
|
||||
|
||||
def read_json(path: Path) -> dict:
|
||||
return json.loads(path.read_text(encoding="utf-8"))
|
||||
|
||||
|
||||
def file_sha256(path: Path) -> str:
|
||||
digest = hashlib.sha256()
|
||||
with path.open("rb") as handle:
|
||||
for chunk in iter(lambda: handle.read(1024 * 1024), b""):
|
||||
digest.update(chunk)
|
||||
return digest.hexdigest()
|
||||
|
||||
|
||||
def evaluate_clinical(
|
||||
project_root: Path,
|
||||
source_manifest: dict,
|
||||
approval_manifest: dict | None = None,
|
||||
) -> List[Gate]:
|
||||
pdf_path = project_root / source_manifest.get("pdf_path", "")
|
||||
expected_sha = source_manifest.get("sha256", "")
|
||||
actual_sha = file_sha256(pdf_path) if pdf_path.is_file() else ""
|
||||
approvals = approval_manifest or {}
|
||||
approved_roles = {
|
||||
item.get("role") for item in approvals.get("approvals", [])
|
||||
if item.get("approved") and item.get("reviewer") and item.get("date")
|
||||
}
|
||||
artifacts = approvals.get("artifacts", {})
|
||||
|
||||
gates = [
|
||||
Gate(
|
||||
"current_national_formulary_edition",
|
||||
int(source_manifest.get("edition") == CURRENT_FORMULARY_EDITION),
|
||||
target=1,
|
||||
detail=(
|
||||
f"found edition {source_manifest.get('edition')}; "
|
||||
f"required edition {CURRENT_FORMULARY_EDITION}"
|
||||
),
|
||||
),
|
||||
Gate(
|
||||
"source_pdf_sha256_matches_manifest",
|
||||
int(bool(expected_sha) and actual_sha == expected_sha),
|
||||
target=1,
|
||||
detail="source PDF missing or hash mismatch" if actual_sha != expected_sha else "",
|
||||
),
|
||||
Gate(
|
||||
"production_use_rights_documented",
|
||||
int(bool(source_manifest.get("production_use_rights_documented"))),
|
||||
target=1,
|
||||
),
|
||||
Gate(
|
||||
"source_marked_clinical_production_eligible",
|
||||
int(bool(source_manifest.get("clinical_production_eligible"))),
|
||||
target=1,
|
||||
),
|
||||
]
|
||||
|
||||
for name, default_path in REQUIRED_RELEASE_ARTIFACTS.items():
|
||||
configured = artifacts.get(name, default_path)
|
||||
path = project_root / configured
|
||||
gates.append(Gate(f"artifact_{name}", int(path.is_file()), target=1,
|
||||
detail=str(configured)))
|
||||
|
||||
missing_roles = sorted(REQUIRED_APPROVAL_ROLES - approved_roles)
|
||||
gates.append(Gate(
|
||||
"required_clinical_release_approvals",
|
||||
len(approved_roles & REQUIRED_APPROVAL_ROLES),
|
||||
target=len(REQUIRED_APPROVAL_ROLES),
|
||||
detail="missing: " + ", ".join(missing_roles) if missing_roles else "",
|
||||
))
|
||||
gates.append(Gate(
|
||||
"intended_use_and_regulatory_classification_approved",
|
||||
int(bool(approvals.get("intended_use_approved")) and
|
||||
bool(approvals.get("regulatory_classification_approved"))),
|
||||
target=1,
|
||||
))
|
||||
return gates
|
||||
@@ -17,6 +17,13 @@ from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Dict, Iterable, List, Sequence
|
||||
|
||||
import re as _re
|
||||
|
||||
_WHITESPACE = _re.compile(r"\s+")
|
||||
# Head of a source line, enough to identify it without demanding that an
|
||||
# intentionally comma-split list match in full.
|
||||
SENTENCE_PROBE_CHARS = 60
|
||||
|
||||
PUA_RANGE = (0xE000, 0xF8FF)
|
||||
REPLACEMENT_CHAR = "�"
|
||||
|
||||
@@ -66,11 +73,12 @@ def _count_pua(text: str) -> int:
|
||||
def evaluate(monographs: Sequence[dict],
|
||||
transcribed_runs: Sequence[dict] = ()) -> List[Gate]:
|
||||
"""Compute every readiness gate over the whole corpus."""
|
||||
pua = replacement = empty = no_provenance = 0
|
||||
pua = replacement = empty = no_provenance = part_no_provenance = 0
|
||||
corruptions: Dict[str, int] = {c: 0 for c in KNOWN_CORRUPTIONS}
|
||||
formula_leaks: Dict[str, int] = {f: 0 for f in FORMULA_FRAGMENTS}
|
||||
unflagged_blocks = 0
|
||||
ids: Dict[str, int] = {}
|
||||
table_ids: Dict[str, int] = {}
|
||||
no_page_range = 0
|
||||
corpus = []
|
||||
|
||||
@@ -85,6 +93,9 @@ def evaluate(monographs: Sequence[dict],
|
||||
empty += 1
|
||||
if not section.get("parts"):
|
||||
no_provenance += 1
|
||||
for part in section.get("parts") or []:
|
||||
if not part.get("source_span_ids"):
|
||||
part_no_provenance += 1
|
||||
pua += _count_pua(text)
|
||||
replacement += text.count(REPLACEMENT_CHAR)
|
||||
for phrase in KNOWN_CORRUPTIONS:
|
||||
@@ -92,6 +103,9 @@ def evaluate(monographs: Sequence[dict],
|
||||
for phrase in FORMULA_FRAGMENTS:
|
||||
formula_leaks[phrase] += text.count(phrase)
|
||||
for block in monograph.get("tables") or []:
|
||||
table_id = block.get("table_id")
|
||||
if table_id:
|
||||
table_ids[table_id] = table_ids.get(table_id, 0) + 1
|
||||
if not block.get("quarantined"):
|
||||
unflagged_blocks += 1
|
||||
|
||||
@@ -113,7 +127,9 @@ def evaluate(monographs: Sequence[dict],
|
||||
Gate("replacement_char_ufffd", replacement),
|
||||
Gate("empty_section", empty),
|
||||
Gate("section_without_provenance", no_provenance),
|
||||
Gate("part_without_source_span_ids", part_no_provenance),
|
||||
Gate("unflagged_quarantine_block", unflagged_blocks),
|
||||
Gate("duplicate_table_id", sum(1 for n in table_ids.values() if n > 1)),
|
||||
Gate("duplicate_drug_id", sum(1 for n in ids.values() if n > 1)),
|
||||
Gate("monograph_without_page_range", no_page_range),
|
||||
]
|
||||
@@ -147,7 +163,10 @@ def evaluate_chunks(monographs: Sequence[dict],
|
||||
blocks_by_section: Dict[tuple, list] = {}
|
||||
block_ids: Dict[str, str] = {}
|
||||
block_texts: Dict[str, str] = {}
|
||||
sections_by_key: Dict[tuple, dict] = {}
|
||||
for monograph in monographs:
|
||||
for section_key, section in (monograph.get("sections") or {}).items():
|
||||
sections_by_key[(monograph["drug_id"], section_key)] = section
|
||||
for block in monograph.get("tables") or []:
|
||||
key = (monograph["drug_id"], block.get("section_key"))
|
||||
blocks_by_section.setdefault(key, []).append(block)
|
||||
@@ -159,15 +178,76 @@ def evaluate_chunks(monographs: Sequence[dict],
|
||||
unknown_id = missing_provenance = leaked = 0
|
||||
descriptors = 0
|
||||
descriptor_without_attachment = 0
|
||||
attachment_header_row_present = 0
|
||||
descriptor_with_unverified_columns = 0
|
||||
unsupported_schema = 0
|
||||
prose_without_source_text = 0
|
||||
source_text_not_unique = 0
|
||||
physical_range_not_exact = 0
|
||||
descriptor_range_not_attachment_page = 0
|
||||
attachment_without_printed_page = 0
|
||||
context_label_missing_from_text = 0
|
||||
invalid_printed_page_range = 0
|
||||
|
||||
for chunk in chunks:
|
||||
if chunk.get("schema_version") != 4:
|
||||
unsupported_schema += 1
|
||||
printed_range = chunk.get("printed_page_range")
|
||||
if (
|
||||
not isinstance(printed_range, list)
|
||||
or len(printed_range) != 2
|
||||
or not all(isinstance(page, int) for page in printed_range)
|
||||
or printed_range[0] > printed_range[1]
|
||||
):
|
||||
invalid_printed_page_range += 1
|
||||
attachments = chunk.get("attachments") or []
|
||||
if chunk.get("chunk_kind") == "block_descriptor":
|
||||
descriptors += 1
|
||||
if not attachments:
|
||||
descriptor_without_attachment += 1
|
||||
if "Cột:" in (chunk.get("text") or ""):
|
||||
descriptor_with_unverified_columns += 1
|
||||
if attachments:
|
||||
page = attachments[0].get("physical_page")
|
||||
if chunk.get("source_page_range") != [page, page]:
|
||||
descriptor_range_not_attachment_page += 1
|
||||
else:
|
||||
source_body = chunk.get("source_text") or ""
|
||||
if not source_body:
|
||||
prose_without_source_text += 1
|
||||
section = sections_by_key.get((chunk["drug_id"], chunk["section_key"]))
|
||||
section_text = (section or {}).get("text", "").strip()
|
||||
if not source_body or section_text.count(source_body) != 1:
|
||||
source_text_not_unique += 1
|
||||
else:
|
||||
chunk_start = section_text.index(source_body)
|
||||
chunk_end = chunk_start + len(source_body)
|
||||
cursor = 0
|
||||
pages = []
|
||||
for part in (section or {}).get("parts") or []:
|
||||
if (part.get("kind") != "prose" or part.get("quarantined")
|
||||
or not part.get("text")):
|
||||
continue
|
||||
part_start = section_text.find(part["text"], cursor)
|
||||
if part_start < 0:
|
||||
pages = []
|
||||
break
|
||||
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"])
|
||||
expected = [min(pages), max(pages)] if pages else None
|
||||
if chunk.get("source_page_range") != expected:
|
||||
physical_range_not_exact += 1
|
||||
body = chunk.get("text") or ""
|
||||
if any(label not in body for label in chunk.get("context_labels") or []):
|
||||
context_label_missing_from_text += 1
|
||||
key = (chunk["drug_id"], chunk["section_key"])
|
||||
for attachment in attachments:
|
||||
if attachment.get("header_row"):
|
||||
attachment_header_row_present += 1
|
||||
if attachment.get("printed_page") is None:
|
||||
attachment_without_printed_page += 1
|
||||
referenced.setdefault(key, set()).add(attachment["block_id"])
|
||||
if block_ids.get(attachment["block_id"]) != chunk["drug_id"]:
|
||||
unknown_id += 1
|
||||
@@ -180,6 +260,9 @@ def evaluate_chunks(monographs: Sequence[dict],
|
||||
if len(probe) > 20 and probe in body:
|
||||
leaked += 1
|
||||
|
||||
over_ceiling = [c for c in chunks if c.get("oversized")]
|
||||
uncovered = _sections_not_covered(monographs, chunks)
|
||||
|
||||
unreferenced = 0
|
||||
for key, blocks in blocks_by_section.items():
|
||||
seen = referenced.get(key, set())
|
||||
@@ -187,10 +270,25 @@ def evaluate_chunks(monographs: Sequence[dict],
|
||||
|
||||
total_blocks = sum(len(v) for v in blocks_by_section.values())
|
||||
return [
|
||||
Gate("chunk_over_token_ceiling", len(over_ceiling),
|
||||
detail="; ".join(c["chunk_id"] for c in over_ceiling[:3])),
|
||||
Gate("chunk_without_printed_page_range", invalid_printed_page_range),
|
||||
Gate("chunk_schema_version_not_supported", unsupported_schema),
|
||||
Gate("prose_without_source_text", prose_without_source_text),
|
||||
Gate("chunk_source_text_not_unique", source_text_not_unique),
|
||||
Gate("chunk_physical_range_not_exact", physical_range_not_exact),
|
||||
Gate("descriptor_range_not_attachment_page",
|
||||
descriptor_range_not_attachment_page),
|
||||
Gate("attachment_without_printed_page", attachment_without_printed_page),
|
||||
Gate("context_label_missing_from_text", context_label_missing_from_text),
|
||||
Gate("section_not_reassemblable_from_chunks", len(uncovered),
|
||||
detail="; ".join(uncovered[:3])),
|
||||
Gate("section_block_without_chunk_reference", unreferenced),
|
||||
Gate("attachment_block_id_unknown", unknown_id),
|
||||
Gate("attachment_without_page_or_bbox", missing_provenance),
|
||||
Gate("block_text_leaked_into_chunk_text", leaked),
|
||||
Gate("attachment_header_row_present", attachment_header_row_present),
|
||||
Gate("descriptor_with_unverified_columns", descriptor_with_unverified_columns),
|
||||
Gate("descriptor_chunk_without_attachment", descriptor_without_attachment),
|
||||
Gate("descriptor_count_vs_block_count", descriptors, target=total_blocks,
|
||||
detail=f"{descriptors} descriptors for {total_blocks} blocks"),
|
||||
@@ -200,3 +298,65 @@ def evaluate_chunks(monographs: Sequence[dict],
|
||||
def read_chunks(path: Path) -> List[dict]:
|
||||
with path.open(encoding="utf-8") as handle:
|
||||
return [json.loads(line) for line in handle if line.strip()]
|
||||
|
||||
|
||||
def _reassemble(parts: Sequence[str]) -> str:
|
||||
"""Glue chunk parts back together, removing the deliberate overlap."""
|
||||
if not parts:
|
||||
return ""
|
||||
text = parts[0]
|
||||
for part in parts[1:]:
|
||||
overlap = 0
|
||||
for size in range(min(len(text), len(part)), 0, -1):
|
||||
if text.endswith(part[:size]):
|
||||
overlap = size
|
||||
break
|
||||
text += part[overlap:]
|
||||
return text
|
||||
|
||||
|
||||
def _sections_not_covered(monographs: Sequence[dict],
|
||||
chunks: Sequence[dict]) -> List[str]:
|
||||
"""Sections that cannot be rebuilt exactly from their own chunks.
|
||||
|
||||
Stronger than asking whether each line still appears somewhere: it proves
|
||||
the chunks are a faithful partition of the section, so nothing was dropped
|
||||
*and* nothing was reordered or duplicated beyond the intended overlap.
|
||||
|
||||
Whitespace is removed from both sides rather than normalised, because
|
||||
each split seam legitimately loses one space: a part is built with
|
||||
`"".join(...).strip()`, and the space that sat between two sentences falls
|
||||
on the boundary. Measured on ABACAVIR: exactly two single spaces across a
|
||||
4,232-character section, nothing else. The gate exists to prove no
|
||||
character of *content* is lost, reordered, or duplicated beyond the
|
||||
intended overlap; it is not a formatting check.
|
||||
|
||||
Two weaker versions were tried first and both were instrument bugs, not
|
||||
data bugs. Joining chunk texts with a newline meant a paragraph split
|
||||
across parts could never match, reporting 734 sections missing when the
|
||||
first one it named was present. Probing a 60-character head then failed on
|
||||
NAPROXEN alone, because that probe straddled an overlap seam where the
|
||||
repeated text legitimately appears twice.
|
||||
"""
|
||||
by_section: Dict[tuple, List[dict]] = {}
|
||||
for chunk in chunks:
|
||||
if chunk.get("chunk_kind") == "block_descriptor":
|
||||
continue
|
||||
key = (chunk["drug_id"], chunk["section_key"])
|
||||
by_section.setdefault(key, []).append(chunk)
|
||||
|
||||
broken = []
|
||||
for monograph in monographs:
|
||||
for key, section in (monograph.get("sections") or {}).items():
|
||||
source = _WHITESPACE.sub("", section.get("text") or "")
|
||||
if not source:
|
||||
continue
|
||||
parts = sorted(by_section.get((monograph["drug_id"], key), []),
|
||||
key=lambda c: c.get("part_index", 0))
|
||||
rebuilt = _WHITESPACE.sub(
|
||||
"", _reassemble([
|
||||
c.get("source_text") or c.get("text") or "" for c in parts
|
||||
]))
|
||||
if rebuilt != source:
|
||||
broken.append(f"{monograph['drug_id']}/{key}")
|
||||
return broken
|
||||
|
||||
@@ -3,10 +3,22 @@ name = "ingestion"
|
||||
version = "0.0.0"
|
||||
description = "Offline batch pipeline: PDF -> monographs -> chunks -> embeddings -> Qdrant"
|
||||
requires-python = ">=3.11"
|
||||
dependencies = ["pymupdf>=1.24", "numpy>=1.26", "scipy>=1.11"]
|
||||
dependencies = ["pymupdf>=1.24", "numpy>=1.26", "scipy>=1.11", "tiktoken>=0.7"]
|
||||
|
||||
[project.optional-dependencies]
|
||||
dev = ["pytest>=7.4"]
|
||||
# Only the live probe and a real embedding run need these. The adapters and
|
||||
# their tests import neither, so the default install stays offline-capable.
|
||||
bedrock = ["boto3>=1.34"]
|
||||
local-embed = ["sentence-transformers>=3.0"]
|
||||
# `load/` talks to a VectorStore port; only `load.qdrant_repo` imports this,
|
||||
# lazily, so the load stage is tested in full with no server running.
|
||||
qdrant = ["qdrant-client>=1.7"]
|
||||
|
||||
[tool.pytest.ini_options]
|
||||
markers = [
|
||||
"integration: needs a live service (a local Qdrant); skips when absent",
|
||||
]
|
||||
|
||||
[build-system]
|
||||
requires = ["setuptools>=68"]
|
||||
|
||||
@@ -6,11 +6,19 @@ from ingestion.chunk import (
|
||||
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 Heading, Monograph, SectionSpan, TableBlock
|
||||
from ingestion.segment.models import (
|
||||
PART_PROSE,
|
||||
Heading,
|
||||
Monograph,
|
||||
SectionPart,
|
||||
SectionSpan,
|
||||
TableBlock,
|
||||
)
|
||||
from ingestion.tables import SHAPE_FORMULA_2D, SHAPE_MULTI_HEADER, SHAPE_SIMPLE
|
||||
|
||||
|
||||
@@ -20,6 +28,15 @@ def _section(key, display, text, page=202):
|
||||
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 []),
|
||||
)
|
||||
|
||||
|
||||
@@ -65,13 +82,17 @@ def test_a_section_whose_table_was_lifted_says_so():
|
||||
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)
|
||||
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():
|
||||
@@ -106,15 +127,17 @@ def test_a_header_row_carrying_a_number_is_refused():
|
||||
assert descriptor.attachments[0].header_row == []
|
||||
|
||||
|
||||
def test_only_a_simple_table_contributes_a_header():
|
||||
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": ["Nhóm", "Liều"]}
|
||||
for shape, expected in ((SHAPE_SIMPLE, ["Nhóm", "Liều"]),
|
||||
(SHAPE_MULTI_HEADER, [])):
|
||||
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 == expected
|
||||
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():
|
||||
@@ -163,6 +186,242 @@ def test_describe_block_names_the_page_even_with_no_header():
|
||||
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)
|
||||
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)
|
||||
|
||||
@@ -0,0 +1,233 @@
|
||||
"""The embedding cache, exercised with a counting stub and no network.
|
||||
|
||||
The claim these tests exist to make checkable is narrow and financial: running
|
||||
the corpus a second time must cost nothing. That is asserted by counting calls
|
||||
the *inner* provider received, not by trusting a hit counter.
|
||||
|
||||
The other half is the inverse — the cases where a hit would be wrong. Serving a
|
||||
vector after its text was edited, across two models, or across Cohere's
|
||||
document/query subspaces would each be silent: no error, just worse recall or a
|
||||
corpus of mixed vectors. There is a test per direction.
|
||||
"""
|
||||
import json
|
||||
|
||||
import pytest
|
||||
|
||||
from ingestion.embed import INPUT_DOCUMENT, EmbeddingVector
|
||||
from ingestion.embed.cache import (
|
||||
CachingEmbeddingProvider,
|
||||
EmbeddingCache,
|
||||
cache_key,
|
||||
)
|
||||
from ingestion.embed.ports import EmbeddingProvider
|
||||
|
||||
DIMENSIONS = 8
|
||||
|
||||
|
||||
class CountingProvider(EmbeddingProvider):
|
||||
"""Deterministic vectors, and a record of every text it was asked for."""
|
||||
|
||||
def __init__(self, model_id="stub-model-v1", dimensions=DIMENSIONS, batch=96):
|
||||
self._model_id = model_id
|
||||
self._dimensions = dimensions
|
||||
self._batch = batch
|
||||
self.embedded_texts = []
|
||||
self.batch_calls = 0
|
||||
|
||||
@property
|
||||
def name(self):
|
||||
return "stub"
|
||||
|
||||
@property
|
||||
def model_id(self):
|
||||
return self._model_id
|
||||
|
||||
@property
|
||||
def dimensions(self):
|
||||
return self._dimensions
|
||||
|
||||
@property
|
||||
def max_batch_size(self):
|
||||
return self._batch
|
||||
|
||||
def _embed_batch(self, texts, input_kind):
|
||||
self.batch_calls += 1
|
||||
self.embedded_texts.extend(texts)
|
||||
return [self._vector(text, input_kind) for text in texts]
|
||||
|
||||
def _vector(self, text, input_kind):
|
||||
from ingestion.embed import text_digest
|
||||
|
||||
seed = len(text) + (0 if input_kind == INPUT_DOCUMENT else 1000)
|
||||
return EmbeddingVector(
|
||||
values=[float(seed + i) for i in range(self._dimensions)],
|
||||
text_sha256=text_digest(text),
|
||||
provider="stub",
|
||||
model_id=self._model_id,
|
||||
dimensions=self._dimensions,
|
||||
input_kind=input_kind,
|
||||
normalized=True,
|
||||
input_token_count=len(text.split()),
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def cache_path(tmp_path):
|
||||
return tmp_path / "embeddings.jsonl"
|
||||
|
||||
|
||||
def test_second_run_over_the_same_texts_makes_zero_provider_requests(cache_path):
|
||||
texts = ["paracetamol", "chống chỉ định", "liều dùng cho trẻ em"]
|
||||
inner = CountingProvider()
|
||||
|
||||
first = CachingEmbeddingProvider(inner, EmbeddingCache(cache_path))
|
||||
cold = first.embed_documents(texts)
|
||||
assert cold.request_count == 1
|
||||
assert inner.embedded_texts == texts
|
||||
|
||||
reopened = EmbeddingCache(cache_path)
|
||||
warm = CachingEmbeddingProvider(inner, reopened).embed_documents(texts)
|
||||
|
||||
assert warm.request_count == 0
|
||||
assert inner.embedded_texts == texts, "no text reached the provider twice"
|
||||
assert reopened.stats.hits == len(texts)
|
||||
assert reopened.stats.misses == 0
|
||||
assert reopened.stats.hit_rate == 1.0
|
||||
assert [v.values for v in warm.vectors] == [v.values for v in cold.vectors]
|
||||
|
||||
|
||||
def test_a_repeated_text_in_one_call_is_embedded_once(cache_path):
|
||||
inner = CountingProvider()
|
||||
provider = CachingEmbeddingProvider(inner, EmbeddingCache(cache_path))
|
||||
|
||||
batch = provider.embed_documents(["Abacavir.", "Abacavir.", "Abacavir."])
|
||||
|
||||
assert inner.embedded_texts == ["Abacavir."]
|
||||
assert len(batch.vectors) == 3
|
||||
assert batch.vectors[0].values == batch.vectors[2].values
|
||||
|
||||
|
||||
def test_editing_the_text_is_a_miss_not_a_stale_hit(cache_path):
|
||||
inner = CountingProvider()
|
||||
cache = EmbeddingCache(cache_path)
|
||||
CachingEmbeddingProvider(inner, cache).embed_documents(["liều 500 mg"])
|
||||
|
||||
CachingEmbeddingProvider(inner, cache).embed_documents(["liều 250 mg"])
|
||||
|
||||
assert inner.embedded_texts == ["liều 500 mg", "liều 250 mg"]
|
||||
|
||||
|
||||
def test_a_second_model_never_reuses_the_first_models_vectors(cache_path):
|
||||
cache = EmbeddingCache(cache_path)
|
||||
titan = CountingProvider(model_id="amazon.titan-embed-text-v2:0")
|
||||
cohere = CountingProvider(model_id="cohere.embed-v4:0")
|
||||
|
||||
CachingEmbeddingProvider(titan, cache).embed_documents(["metformin"])
|
||||
CachingEmbeddingProvider(cohere, cache).embed_documents(["metformin"])
|
||||
|
||||
assert titan.embedded_texts == ["metformin"]
|
||||
assert cohere.embedded_texts == ["metformin"]
|
||||
assert len(cache) == 2
|
||||
|
||||
|
||||
def test_query_and_document_kinds_are_cached_separately(cache_path):
|
||||
inner = CountingProvider()
|
||||
cache = EmbeddingCache(cache_path)
|
||||
provider = CachingEmbeddingProvider(inner, cache)
|
||||
|
||||
as_document = provider.embed_documents(["aspirin"])
|
||||
as_query = provider.embed_queries(["aspirin"])
|
||||
|
||||
assert inner.embedded_texts == ["aspirin", "aspirin"]
|
||||
assert as_document.vectors[0].values != as_query.vectors[0].values
|
||||
assert len(cache) == 2
|
||||
|
||||
|
||||
def test_index_and_values_survive_reopening_the_file(cache_path):
|
||||
inner = CountingProvider()
|
||||
original = CachingEmbeddingProvider(
|
||||
inner, EmbeddingCache(cache_path)
|
||||
).embed_documents(["ACETAZOLAMID", "ADENOSIN"])
|
||||
|
||||
reopened = EmbeddingCache(cache_path)
|
||||
|
||||
assert len(reopened) == 2
|
||||
restored = reopened.get(cache_key(inner.model_id, INPUT_DOCUMENT, "ADENOSIN"))
|
||||
assert restored is not None
|
||||
assert restored.values == original.vectors[1].values
|
||||
assert restored.input_kind == INPUT_DOCUMENT
|
||||
assert restored.normalized is True
|
||||
assert restored.input_token_count == 1
|
||||
|
||||
|
||||
def test_putting_a_key_twice_does_not_append_a_second_record(cache_path):
|
||||
cache = EmbeddingCache(cache_path)
|
||||
inner = CountingProvider()
|
||||
vector = inner.embed_documents(["digoxin"]).vectors[0]
|
||||
|
||||
assert cache.put(vector) is True
|
||||
assert cache.put(vector) is False
|
||||
|
||||
lines = cache_path.read_text(encoding="utf-8").strip().splitlines()
|
||||
assert len(lines) == 1
|
||||
assert len(cache) == 1
|
||||
|
||||
|
||||
def test_a_cached_record_whose_length_contradicts_its_dimensions_is_rejected(
|
||||
cache_path,
|
||||
):
|
||||
record = {
|
||||
"model_id": "stub-model-v1",
|
||||
"input_kind": INPUT_DOCUMENT,
|
||||
"text_sha256": "a" * 64,
|
||||
"provider": "stub",
|
||||
"dimensions": 1024,
|
||||
"normalized": True,
|
||||
"input_token_count": 3,
|
||||
"values": [0.1, 0.2],
|
||||
}
|
||||
cache_path.write_text(json.dumps(record) + "\n", encoding="utf-8")
|
||||
cache = EmbeddingCache(cache_path)
|
||||
|
||||
with pytest.raises(ValueError, match="declares 1024 dimensions"):
|
||||
cache.get(("stub-model-v1", INPUT_DOCUMENT, "a" * 64))
|
||||
|
||||
|
||||
def test_a_record_missing_key_fields_is_rejected_at_index_time(cache_path):
|
||||
cache_path.write_text(
|
||||
json.dumps({"model_id": "stub", "values": []}) + "\n", encoding="utf-8"
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="missing key fields"):
|
||||
EmbeddingCache(cache_path)
|
||||
|
||||
|
||||
def test_the_cache_wrapper_still_refuses_empty_text_and_bad_input_kind(cache_path):
|
||||
provider = CachingEmbeddingProvider(CountingProvider(), EmbeddingCache(cache_path))
|
||||
|
||||
with pytest.raises(ValueError, match="empty or whitespace-only"):
|
||||
provider.embed_documents(["paracetamol", " "])
|
||||
with pytest.raises(ValueError, match="input_kind must be one of"):
|
||||
provider.embed(["paracetamol"], "search_document")
|
||||
|
||||
|
||||
def test_a_missing_cache_file_starts_empty_and_is_created_on_first_put(cache_path):
|
||||
cache = EmbeddingCache(cache_path)
|
||||
|
||||
assert len(cache) == 0
|
||||
assert not cache_path.exists()
|
||||
|
||||
CachingEmbeddingProvider(CountingProvider(), cache).embed_documents(["insulin"])
|
||||
|
||||
assert cache_path.exists()
|
||||
assert len(cache) == 1
|
||||
|
||||
|
||||
def test_wrapper_reports_the_inner_models_identity_not_its_own(cache_path):
|
||||
inner = CountingProvider(model_id="cohere.embed-v4:0", dimensions=DIMENSIONS)
|
||||
provider = CachingEmbeddingProvider(inner, EmbeddingCache(cache_path))
|
||||
|
||||
assert provider.model_id == "cohere.embed-v4:0"
|
||||
assert provider.dimensions == DIMENSIONS
|
||||
assert provider.max_batch_size == inner.max_batch_size
|
||||
assert provider.name == "cached:stub"
|
||||
@@ -0,0 +1,324 @@
|
||||
"""Provider adapters, exercised with no AWS account and no network.
|
||||
|
||||
Every Bedrock call goes through a recording stub, so what is under test is the
|
||||
part that can actually be wrong offline: the request body we send, and our
|
||||
reading of the response bodies AWS documents. The one thing these tests cannot
|
||||
establish is whether AWS accepts that body — that needs the live probe, and
|
||||
the coordination handoff says so explicitly.
|
||||
"""
|
||||
import io
|
||||
import json
|
||||
import math
|
||||
|
||||
import pytest
|
||||
|
||||
from ingestion.embed import (
|
||||
BGE_M3,
|
||||
COHERE_V4,
|
||||
TITAN_V2,
|
||||
INPUT_DOCUMENT,
|
||||
INPUT_QUERY,
|
||||
Boto3BedrockInvoker,
|
||||
EmbeddingVector,
|
||||
build_provider,
|
||||
provider_names,
|
||||
text_digest,
|
||||
)
|
||||
from ingestion.embed import probe
|
||||
from ingestion.embed.bedrock_cohere import CohereEmbedV4
|
||||
from ingestion.embed.bedrock_titan import TitanTextEmbeddingsV2
|
||||
from ingestion.embed.local_bge_m3 import BgeM3Local
|
||||
|
||||
|
||||
class RecordingInvoker:
|
||||
"""Stands in for Bedrock; remembers every request it was handed."""
|
||||
|
||||
def __init__(self, responses):
|
||||
self._responses = list(responses)
|
||||
self.calls = []
|
||||
|
||||
def invoke_json(self, model_id, payload, accept="application/json"):
|
||||
self.calls.append(
|
||||
{"model_id": model_id, "payload": payload, "accept": accept}
|
||||
)
|
||||
return self._responses.pop(0)
|
||||
|
||||
|
||||
def _titan_response(dimensions=1024, token_count=12):
|
||||
return {
|
||||
"embedding": [0.01] * dimensions,
|
||||
"inputTextTokenCount": token_count,
|
||||
"embeddingsByType": {"float": [0.01] * dimensions},
|
||||
}
|
||||
|
||||
|
||||
def _cohere_by_type_response(rows, dimensions=1024):
|
||||
return {
|
||||
"id": "stub-id",
|
||||
"response_type": "embeddings_by_type",
|
||||
"embeddings": {"float": [[0.02] * dimensions for _ in range(rows)]},
|
||||
"texts": ["stub"] * rows,
|
||||
}
|
||||
|
||||
|
||||
def _cohere_floats_response(rows, dimensions=1024):
|
||||
return {
|
||||
"id": "stub-id",
|
||||
"response_type": "embeddings_floats",
|
||||
"embeddings": [[0.02] * dimensions for _ in range(rows)],
|
||||
}
|
||||
|
||||
|
||||
def test_titan_request_body_matches_the_documented_v2_shape():
|
||||
invoker = RecordingInvoker([_titan_response()])
|
||||
provider = TitanTextEmbeddingsV2(invoker, dimensions=1024, normalize=True)
|
||||
|
||||
provider.embed_documents(["paracetamol"])
|
||||
|
||||
payload = invoker.calls[0]["payload"]
|
||||
assert invoker.calls[0]["model_id"] == "amazon.titan-embed-text-v2:0"
|
||||
assert payload == {
|
||||
"inputText": "paracetamol",
|
||||
"dimensions": 1024,
|
||||
"normalize": True,
|
||||
}
|
||||
|
||||
|
||||
def test_titan_records_provenance_and_reported_token_count():
|
||||
invoker = RecordingInvoker([_titan_response(token_count=7)])
|
||||
provider = TitanTextEmbeddingsV2(invoker)
|
||||
|
||||
vector = provider.embed_documents(["paracetamol"]).vectors[0]
|
||||
|
||||
assert vector.model_id == "amazon.titan-embed-text-v2:0"
|
||||
assert vector.provider == TITAN_V2
|
||||
assert vector.dimensions == 1024
|
||||
assert vector.input_kind == INPUT_DOCUMENT
|
||||
assert vector.normalized is True
|
||||
assert vector.input_token_count == 7
|
||||
assert vector.text_sha256 == text_digest("paracetamol")
|
||||
|
||||
|
||||
def test_titan_sends_one_request_per_text():
|
||||
invoker = RecordingInvoker([_titan_response(), _titan_response()])
|
||||
provider = TitanTextEmbeddingsV2(invoker)
|
||||
|
||||
batch = provider.embed_documents(["a", "b"])
|
||||
|
||||
assert batch.request_count == 2
|
||||
assert len(batch.vectors) == 2
|
||||
|
||||
|
||||
def test_titan_rejects_a_dimension_the_model_does_not_offer():
|
||||
with pytest.raises(ValueError, match="supports"):
|
||||
TitanTextEmbeddingsV2(RecordingInvoker([]), dimensions=768)
|
||||
|
||||
|
||||
def test_cohere_uses_search_document_for_corpus_and_search_query_for_queries():
|
||||
invoker = RecordingInvoker(
|
||||
[_cohere_by_type_response(1), _cohere_by_type_response(1)]
|
||||
)
|
||||
provider = CohereEmbedV4(invoker)
|
||||
|
||||
provider.embed_documents(["metformin"])
|
||||
provider.embed_queries(["liều metformin"])
|
||||
|
||||
assert invoker.calls[0]["payload"]["input_type"] == "search_document"
|
||||
assert invoker.calls[1]["payload"]["input_type"] == "search_query"
|
||||
|
||||
|
||||
def test_cohere_request_body_pins_dimension_float_type_and_no_truncation():
|
||||
invoker = RecordingInvoker([_cohere_by_type_response(2)])
|
||||
provider = CohereEmbedV4(invoker, dimensions=1024)
|
||||
|
||||
provider.embed_documents(["a", "b"])
|
||||
|
||||
payload = invoker.calls[0]["payload"]
|
||||
assert invoker.calls[0]["model_id"] == "cohere.embed-v4:0"
|
||||
assert payload["texts"] == ["a", "b"]
|
||||
assert payload["embedding_types"] == ["float"]
|
||||
# Left unset the model would return 1536, which no 1024-wide collection
|
||||
# can accept.
|
||||
assert payload["output_dimension"] == 1024
|
||||
# An over-length input must fail, not arrive silently shortened.
|
||||
assert payload["truncate"] == "NONE"
|
||||
assert invoker.calls[0]["accept"] == "*/*"
|
||||
|
||||
|
||||
def test_cohere_reads_the_embeddings_by_type_response():
|
||||
invoker = RecordingInvoker([_cohere_by_type_response(2)])
|
||||
provider = CohereEmbedV4(invoker)
|
||||
|
||||
batch = provider.embed_documents(["a", "b"])
|
||||
|
||||
assert len(batch.vectors) == 2
|
||||
assert all(len(v.values) == 1024 for v in batch.vectors)
|
||||
assert batch.request_count == 1
|
||||
|
||||
|
||||
def test_cohere_also_reads_the_plain_embeddings_floats_response():
|
||||
invoker = RecordingInvoker([_cohere_floats_response(2)])
|
||||
provider = CohereEmbedV4(invoker)
|
||||
|
||||
batch = provider.embed_documents(["a", "b"])
|
||||
|
||||
assert len(batch.vectors) == 2
|
||||
assert all(len(v.values) == 1024 for v in batch.vectors)
|
||||
|
||||
|
||||
def test_cohere_leaves_normalization_unknown_because_the_docs_do_not_say():
|
||||
invoker = RecordingInvoker([_cohere_by_type_response(1)])
|
||||
|
||||
vector = CohereEmbedV4(invoker).embed_documents(["a"]).vectors[0]
|
||||
|
||||
assert vector.normalized is None
|
||||
|
||||
|
||||
def test_cohere_splits_at_the_documented_96_text_ceiling():
|
||||
invoker = RecordingInvoker(
|
||||
[_cohere_by_type_response(96), _cohere_by_type_response(4)]
|
||||
)
|
||||
provider = CohereEmbedV4(invoker)
|
||||
|
||||
batch = provider.embed_documents([f"t{i}" for i in range(100)])
|
||||
|
||||
assert batch.request_count == 2
|
||||
assert len(invoker.calls[0]["payload"]["texts"]) == 96
|
||||
assert len(invoker.calls[1]["payload"]["texts"]) == 4
|
||||
assert len(batch.vectors) == 100
|
||||
|
||||
|
||||
def test_cohere_rejects_a_batch_size_above_the_documented_ceiling():
|
||||
with pytest.raises(ValueError, match="batch_size"):
|
||||
CohereEmbedV4(RecordingInvoker([]), batch_size=97)
|
||||
|
||||
|
||||
def test_a_wrong_width_vector_fails_instead_of_entering_the_corpus():
|
||||
invoker = RecordingInvoker([_titan_response(dimensions=512)])
|
||||
provider = TitanTextEmbeddingsV2(invoker, dimensions=1024)
|
||||
|
||||
with pytest.raises(ValueError, match="512 dimensions"):
|
||||
provider.embed_documents(["a"])
|
||||
|
||||
|
||||
def test_a_response_missing_its_vectors_fails_loudly():
|
||||
invoker = RecordingInvoker([{"id": "stub", "response_type": "x"}])
|
||||
|
||||
with pytest.raises(ValueError, match="no 'embeddings' field"):
|
||||
CohereEmbedV4(invoker).embed_documents(["a"])
|
||||
|
||||
|
||||
def test_a_count_mismatch_between_texts_and_vectors_fails():
|
||||
invoker = RecordingInvoker([_cohere_by_type_response(1)])
|
||||
|
||||
with pytest.raises(ValueError, match="1 vectors for 2 texts"):
|
||||
CohereEmbedV4(invoker).embed_documents(["a", "b"])
|
||||
|
||||
|
||||
def test_an_unknown_input_kind_is_refused_before_any_request_is_made():
|
||||
invoker = RecordingInvoker([])
|
||||
|
||||
with pytest.raises(ValueError, match="input_kind"):
|
||||
CohereEmbedV4(invoker).embed(["a"], "search_document")
|
||||
assert invoker.calls == []
|
||||
|
||||
|
||||
def test_empty_text_is_refused_before_any_request_is_made():
|
||||
invoker = RecordingInvoker([])
|
||||
|
||||
with pytest.raises(ValueError, match="empty"):
|
||||
CohereEmbedV4(invoker).embed_documents(["a", " "])
|
||||
assert invoker.calls == []
|
||||
|
||||
|
||||
def test_bge_m3_runs_through_an_injected_encoder_with_no_weights_loaded():
|
||||
seen = []
|
||||
|
||||
def encoder(texts):
|
||||
seen.append(list(texts))
|
||||
unit = 1.0 / math.sqrt(1024)
|
||||
return [[unit] * 1024 for _ in texts]
|
||||
|
||||
provider = BgeM3Local(encoder=encoder, batch_size=2)
|
||||
batch = provider.embed_queries(["a", "b", "c"])
|
||||
|
||||
assert seen == [["a", "b"], ["c"]]
|
||||
assert batch.request_count == 2
|
||||
assert len(batch.vectors) == 3
|
||||
assert batch.vectors[0].input_kind == INPUT_QUERY
|
||||
assert batch.vectors[0].model_id == "BAAI/bge-m3"
|
||||
# Injected encoder: we did not set normalize_embeddings, so we do not claim it.
|
||||
assert batch.vectors[0].normalized is None
|
||||
|
||||
|
||||
def test_registry_builds_every_provider_without_touching_an_sdk():
|
||||
assert set(provider_names()) == {TITAN_V2, COHERE_V4, BGE_M3}
|
||||
|
||||
titan = build_provider(TITAN_V2, invoker=RecordingInvoker([]))
|
||||
cohere = build_provider(COHERE_V4, invoker=RecordingInvoker([]))
|
||||
local = build_provider(BGE_M3)
|
||||
|
||||
assert (titan.dimensions, cohere.dimensions, local.dimensions) == (
|
||||
1024,
|
||||
1024,
|
||||
1024,
|
||||
)
|
||||
assert titan.max_batch_size == 1
|
||||
assert cohere.max_batch_size == 96
|
||||
|
||||
|
||||
def test_registry_rejects_an_unknown_provider_name():
|
||||
with pytest.raises(ValueError, match="unknown embedding provider"):
|
||||
build_provider("text-embedding-3-small")
|
||||
|
||||
|
||||
class FakeBotoClient:
|
||||
"""The shape boto3's bedrock-runtime client returns: a streaming body."""
|
||||
|
||||
def __init__(self, response_body):
|
||||
self._response_body = response_body
|
||||
self.kwargs = None
|
||||
|
||||
def invoke_model(self, **kwargs):
|
||||
self.kwargs = kwargs
|
||||
return {"body": io.BytesIO(json.dumps(self._response_body).encode())}
|
||||
|
||||
|
||||
def test_boto3_invoker_serialises_the_request_and_reads_the_streamed_body():
|
||||
client = FakeBotoClient({"embedding": [0.5]})
|
||||
invoker = Boto3BedrockInvoker(region="us-east-1", client=client)
|
||||
|
||||
body = invoker.invoke_json("some.model", {"inputText": "à"}, accept="*/*")
|
||||
|
||||
assert body == {"embedding": [0.5]}
|
||||
assert client.kwargs["modelId"] == "some.model"
|
||||
assert client.kwargs["contentType"] == "application/json"
|
||||
assert client.kwargs["accept"] == "*/*"
|
||||
# Vietnamese must survive the round trip as characters, not \\u escapes
|
||||
# the model would then embed literally.
|
||||
assert json.loads(client.kwargs["body"]) == {"inputText": "à"}
|
||||
|
||||
|
||||
def test_probe_measures_the_l2_norm_rather_than_trusting_the_docs():
|
||||
unit = 1.0 / math.sqrt(4)
|
||||
assert probe._l2_norm([unit] * 4) == pytest.approx(1.0)
|
||||
assert probe._l2_norm([3.0, 4.0]) == pytest.approx(5.0)
|
||||
|
||||
|
||||
def test_probe_reports_a_vector_without_raising(capsys):
|
||||
vector = EmbeddingVector(
|
||||
values=[0.5, 0.5, 0.5, 0.5],
|
||||
text_sha256=text_digest("x"),
|
||||
provider=TITAN_V2,
|
||||
model_id="amazon.titan-embed-text-v2:0",
|
||||
dimensions=4,
|
||||
input_kind=INPUT_DOCUMENT,
|
||||
normalized=True,
|
||||
input_token_count=3,
|
||||
)
|
||||
|
||||
probe._report(vector, latency_ms=12.5, requests=1)
|
||||
|
||||
out = capsys.readouterr().out
|
||||
assert "amazon.titan-embed-text-v2:0" in out
|
||||
assert "measured L2 norm: 1.000000" in out
|
||||
@@ -0,0 +1,40 @@
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from ingestion.entities.catalog import build_entities
|
||||
|
||||
DATA = Path(__file__).parents[1] / "data"
|
||||
PDF = DATA / "raw/duoc-thu-quoc-gia-viet-nam-2018.pdf"
|
||||
MONOGRAPHS = DATA / "processed/monographs.jsonl"
|
||||
|
||||
needs_source = pytest.mark.skipif(
|
||||
not PDF.exists() or not MONOGRAPHS.exists(),
|
||||
reason="source corpus artifacts are not available",
|
||||
)
|
||||
|
||||
|
||||
@needs_source
|
||||
def test_verified_entity_catalog_maps_every_explicit_see_alias():
|
||||
payload = build_entities(MONOGRAPHS, PDF)
|
||||
stats = payload["stats"]
|
||||
assert stats["entity_count"] == 684
|
||||
assert stats["back_index_see_relations"] == 344
|
||||
assert stats["back_index_aliases_mapped"] == 344
|
||||
assert stats["back_index_aliases_unresolved"] == 0
|
||||
assert stats["back_index_aliases_ambiguous"] == 0
|
||||
assert stats["trade_name_sections"] == 492
|
||||
|
||||
|
||||
@needs_source
|
||||
def test_common_parenthesized_names_are_emitted_as_aliases():
|
||||
payload = build_entities(MONOGRAPHS, PDF)
|
||||
entities = {item["drug_id"]: item for item in payload["entities"]}
|
||||
aliases = {
|
||||
drug_id: {alias.casefold() for alias in entity["aliases"]}
|
||||
for drug_id, entity in entities.items()
|
||||
}
|
||||
assert "paracetamol" in aliases["paracetamol_acetaminophen"]
|
||||
assert "acetaminophen" in aliases["paracetamol_acetaminophen"]
|
||||
assert "aspirin" in aliases["acid_acetylsalicylic_aspirin"]
|
||||
assert "oresol" in aliases["thuoc_uong_bu_nuoc_va_ien_giai"]
|
||||
@@ -2,6 +2,7 @@ import json
|
||||
from pathlib import Path
|
||||
|
||||
from ingestion.extract.formulas import (
|
||||
BARLESS_FORMULA_BOTTOM_PT,
|
||||
FORMULA_BAND_HEIGHT_PT,
|
||||
FORMULA_SIDE_MARGIN_PT,
|
||||
load_formula_regions,
|
||||
@@ -50,6 +51,16 @@ def test_the_barless_adenosin_formula_is_recorded_as_a_recall_limit():
|
||||
assert "UNMEASURED" in payload["recall_limit"]
|
||||
|
||||
|
||||
def test_barless_adenosin_band_reaches_its_printed_denominator():
|
||||
payload = json.loads(VERIFIED.read_text(encoding="utf-8"))
|
||||
entry = next(r for r in payload["regions"] if r.get("source_prints_no_bar"))
|
||||
region = next(r for r in load_formula_regions() if r.physical_page == 147)
|
||||
assert region.bbox[3] == entry["bar_bbox"][3] + BARLESS_FORMULA_BOTTOM_PT
|
||||
# "Ví dụ:" begins immediately afterwards with its span centre at ~704.3;
|
||||
# the formula band must stop before that prose and the following table.
|
||||
assert region.bbox[3] < 704
|
||||
|
||||
|
||||
def test_outlined_text_transcriptions_cover_every_detected_run():
|
||||
payload = json.loads(TRANSCRIPTIONS.read_text(encoding="utf-8"))
|
||||
runs = payload["runs"]
|
||||
|
||||
@@ -0,0 +1,589 @@
|
||||
"""The load stage, exercised against the in-memory store with no server.
|
||||
|
||||
Two classes of failure are silent in a vector database and are what most of
|
||||
these tests aim at. Loading the same corpus twice can leave two copies of a
|
||||
dose, and every query still succeeds — so idempotency is asserted by point
|
||||
count, not by inspecting the upsert calls. Mixing two corpus generations or two
|
||||
models into one collection also raises nothing at query time; every search
|
||||
returns *something*, just from the wrong material. That is what the manifest
|
||||
gate exists to make loud, and there is a test per way it can be violated.
|
||||
|
||||
The last test runs over the real `chunks.jsonl` when it is present, because a
|
||||
provenance rule that only holds for hand-written records is not evidence.
|
||||
"""
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from ingestion.load import (
|
||||
ChunkLoader,
|
||||
CollectionSpec,
|
||||
CorpusManifest,
|
||||
CorpusMismatch,
|
||||
InMemoryVectorStore,
|
||||
PointCountMismatch,
|
||||
build_point,
|
||||
corpus_sha256,
|
||||
count_chunks,
|
||||
iter_chunk_records,
|
||||
manifest_collection,
|
||||
point_id_for,
|
||||
read_manifest,
|
||||
validate_chunk_record,
|
||||
)
|
||||
|
||||
DIMENSIONS = 4
|
||||
COLLECTION = "duoc_thu_chunks"
|
||||
CORPUS_SHA = "a" * 64
|
||||
OTHER_SHA = "b" * 64
|
||||
MODEL = "amazon.titan-embed-text-v2:0"
|
||||
|
||||
REAL_CHUNKS = (
|
||||
Path(__file__).resolve().parents[1] / "data" / "processed" / "chunks.jsonl"
|
||||
)
|
||||
|
||||
|
||||
def chunk_record(chunk_id="abacavir__lieu_luong__0", **overrides):
|
||||
record = {
|
||||
"schema_version": 4,
|
||||
"chunk_id": chunk_id,
|
||||
"drug_id": "abacavir",
|
||||
"drug_name": "ABACAVIR",
|
||||
"section_key": "lieu_luong_va_cach_dung",
|
||||
"section_display_name": "Liều lượng và cách dùng",
|
||||
"text": "Người lớn: 300 mg, hai lần mỗi ngày.",
|
||||
"source_text": "Người lớn: 300 mg, hai lần mỗi ngày.",
|
||||
"heading_physical_page": 100,
|
||||
"source_page_range": [100, 102],
|
||||
"printed_page_range": [101, 103],
|
||||
"atc_codes": ["J05AF06"],
|
||||
"part_index": 0,
|
||||
"part_count": 1,
|
||||
"est_tokens": 14,
|
||||
"oversized": False,
|
||||
"chunk_kind": "prose",
|
||||
"attachments": [],
|
||||
"has_quarantined_content": False,
|
||||
}
|
||||
record.update(overrides)
|
||||
return record
|
||||
|
||||
|
||||
def vector(seed=0.1):
|
||||
return [seed] * DIMENSIONS
|
||||
|
||||
|
||||
def manifest(**overrides):
|
||||
values = {
|
||||
"corpus_sha256": CORPUS_SHA,
|
||||
"chunk_count": 3,
|
||||
"model_id": MODEL,
|
||||
"dimensions": DIMENSIONS,
|
||||
"input_kind": "document",
|
||||
"provider": "titan-v2",
|
||||
}
|
||||
values.update(overrides)
|
||||
return CorpusManifest(**values)
|
||||
|
||||
|
||||
def spec(**overrides):
|
||||
values = {"name": COLLECTION, "vector_size": DIMENSIONS}
|
||||
values.update(overrides)
|
||||
return CollectionSpec(**values)
|
||||
|
||||
|
||||
def loader(store, **overrides):
|
||||
return ChunkLoader(
|
||||
store,
|
||||
overrides.pop("spec", spec()),
|
||||
overrides.pop("manifest", manifest()),
|
||||
**overrides,
|
||||
)
|
||||
|
||||
|
||||
def pairs(count=3):
|
||||
return [
|
||||
(chunk_record(chunk_id=f"drug__section__{i}"), vector(0.1 * (i + 1)))
|
||||
for i in range(count)
|
||||
]
|
||||
|
||||
|
||||
# --- A5: derived ids and idempotency -------------------------------------
|
||||
|
||||
|
||||
def test_point_id_is_derived_from_chunk_id_and_is_stable():
|
||||
first = point_id_for("abacavir__lieu_luong__0")
|
||||
second = point_id_for("abacavir__lieu_luong__0")
|
||||
|
||||
assert first == second
|
||||
assert first != point_id_for("abacavir__lieu_luong__1")
|
||||
|
||||
|
||||
def test_point_id_refuses_an_empty_chunk_id():
|
||||
with pytest.raises(ValueError, match="chunk_id is required"):
|
||||
point_id_for(" ")
|
||||
|
||||
|
||||
def test_loading_the_same_corpus_twice_leaves_the_point_count_unchanged():
|
||||
store = InMemoryVectorStore()
|
||||
data = pairs(3)
|
||||
|
||||
first = loader(store).load(data)
|
||||
second = loader(store).load(data)
|
||||
|
||||
assert first.collection_created is True
|
||||
assert second.collection_created is False
|
||||
assert first.collection_count == 3
|
||||
assert second.collection_count == 3, "a re-run duplicated points"
|
||||
assert second.points_upserted == 3
|
||||
|
||||
|
||||
def test_a_reloaded_chunk_overwrites_its_own_point_rather_than_adding_one():
|
||||
store = InMemoryVectorStore()
|
||||
record = chunk_record()
|
||||
loader(store).load([(record, vector(0.1))])
|
||||
|
||||
edited = chunk_record(text="Người lớn: 600 mg, một lần mỗi ngày.")
|
||||
loader(store).load([(edited, vector(0.9))])
|
||||
|
||||
assert store.count(COLLECTION) == 1
|
||||
stored = store.retrieve(COLLECTION, point_id_for(record["chunk_id"]))
|
||||
assert stored.payload["text"] == "Người lớn: 600 mg, một lần mỗi ngày."
|
||||
assert stored.vector == vector(0.9)
|
||||
|
||||
|
||||
def test_records_are_upserted_in_batches_of_the_configured_size():
|
||||
store = InMemoryVectorStore()
|
||||
|
||||
report = loader(store, batch_size=2).load(pairs(5))
|
||||
|
||||
assert report.batches == 3
|
||||
assert report.points_upserted == 5
|
||||
assert report.collection_count == 5
|
||||
|
||||
|
||||
def test_batch_size_must_be_positive():
|
||||
with pytest.raises(ValueError, match="batch_size must be positive"):
|
||||
ChunkLoader(InMemoryVectorStore(), spec(), manifest(), batch_size=0)
|
||||
|
||||
|
||||
# --- A4: collection shape and payload ------------------------------------
|
||||
|
||||
|
||||
def test_collection_is_created_with_the_declared_size_and_payload_indexes():
|
||||
store = InMemoryVectorStore()
|
||||
|
||||
loader(store).load(pairs(1))
|
||||
|
||||
assert store.spec(COLLECTION).vector_size == DIMENSIONS
|
||||
assert store.spec(COLLECTION).distance == "Cosine"
|
||||
indexed = dict(store.indexed_fields(COLLECTION))
|
||||
assert indexed["drug_id"] == "keyword"
|
||||
assert indexed["section_key"] == "keyword"
|
||||
assert indexed["atc_codes"] == "keyword"
|
||||
assert indexed["chunk_kind"] == "keyword"
|
||||
assert indexed["has_quarantined_content"] == "bool"
|
||||
|
||||
|
||||
def test_payload_carries_every_provenance_field_of_the_chunk_record():
|
||||
store = InMemoryVectorStore()
|
||||
record = chunk_record()
|
||||
|
||||
loader(store).load([(record, vector())])
|
||||
|
||||
payload = store.retrieve(COLLECTION, point_id_for(record["chunk_id"])).payload
|
||||
assert payload == record
|
||||
|
||||
|
||||
def test_a_field_added_by_a_future_chunker_flows_through_untouched():
|
||||
store = InMemoryVectorStore()
|
||||
record = chunk_record(
|
||||
population_tags=["Người lớn", "Suy thận"], printed_page_range=[101, 103]
|
||||
)
|
||||
|
||||
loader(store).load([(record, vector())])
|
||||
|
||||
payload = store.retrieve(COLLECTION, point_id_for(record["chunk_id"])).payload
|
||||
assert payload["population_tags"] == ["Người lớn", "Suy thận"]
|
||||
assert payload["printed_page_range"] == [101, 103]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"field",
|
||||
[
|
||||
"chunk_id",
|
||||
"drug_id",
|
||||
"section_key",
|
||||
"source_page_range",
|
||||
"printed_page_range",
|
||||
"text",
|
||||
],
|
||||
)
|
||||
def test_a_chunk_missing_a_required_provenance_field_is_refused(field):
|
||||
with pytest.raises(ValueError, match="missing required provenance fields"):
|
||||
validate_chunk_record(chunk_record(**{field: None}))
|
||||
|
||||
|
||||
# --- failing closed on incomplete provenance ------------------------------
|
||||
#
|
||||
# Every case below passed an earlier version of this validator. The cost of
|
||||
# that is not an exception at load time — it is paying for an embedding run and
|
||||
# then discovering every answer abstains because the chunks cannot be cited.
|
||||
|
||||
|
||||
@pytest.mark.parametrize("field", ["source_page_range", "printed_page_range"])
|
||||
def test_an_empty_page_range_is_missing_not_present(field):
|
||||
"""`[] in (None, "")` is False, which is exactly how this slipped through."""
|
||||
with pytest.raises(ValueError, match="missing required provenance fields"):
|
||||
validate_chunk_record(chunk_record(**{field: []}))
|
||||
|
||||
|
||||
def test_an_unknown_old_or_future_schema_is_refused_fail_closed():
|
||||
for version in (3, 5):
|
||||
with pytest.raises(ValueError, match="supports exactly v4"):
|
||||
validate_chunk_record(chunk_record(schema_version=version))
|
||||
|
||||
|
||||
def test_a_chunk_with_no_schema_version_at_all_is_refused():
|
||||
record = chunk_record()
|
||||
del record["schema_version"]
|
||||
with pytest.raises(ValueError, match="declares schema_version None"):
|
||||
validate_chunk_record(record)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"value", [[101], [101, 102, 103], "101-103", 101, {"start": 101}]
|
||||
)
|
||||
def test_a_page_range_that_is_not_a_pair_is_refused(value):
|
||||
with pytest.raises(ValueError, match=r"expected a \[start, end\] pair"):
|
||||
validate_chunk_record(chunk_record(printed_page_range=value))
|
||||
|
||||
|
||||
def test_a_page_range_running_backwards_is_refused():
|
||||
with pytest.raises(ValueError, match="running backwards"):
|
||||
validate_chunk_record(chunk_record(printed_page_range=[103, 101]))
|
||||
|
||||
|
||||
def test_a_non_integer_page_is_refused():
|
||||
with pytest.raises(ValueError, match="non-integer page"):
|
||||
validate_chunk_record(chunk_record(printed_page_range=[101.5, 103]))
|
||||
|
||||
|
||||
def test_boolean_pages_are_not_accepted_as_python_integers():
|
||||
with pytest.raises(ValueError, match="non-integer page"):
|
||||
validate_chunk_record(chunk_record(printed_page_range=[False, True]))
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("field", "value"),
|
||||
[("source_page_range", [-1, 0]), ("printed_page_range", [0, 1])],
|
||||
)
|
||||
def test_page_ranges_reject_impossible_lower_bounds(field, value):
|
||||
with pytest.raises(ValueError, match="pages must start"):
|
||||
validate_chunk_record(chunk_record(**{field: value}))
|
||||
|
||||
|
||||
def test_page_zero_and_false_are_values_not_absences():
|
||||
"""Physical pages are 0-indexed; a falsiness test would reject real records."""
|
||||
validate_chunk_record(
|
||||
chunk_record(
|
||||
heading_physical_page=0,
|
||||
source_page_range=[0, 0],
|
||||
printed_page_range=[1, 1],
|
||||
has_quarantined_content=False,
|
||||
oversized=False,
|
||||
part_index=0,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def test_a_wrong_sized_vector_is_refused_before_anything_is_upserted():
|
||||
store = InMemoryVectorStore()
|
||||
|
||||
with pytest.raises(ValueError, match="expects 4"):
|
||||
loader(store).load([(chunk_record(), [0.1, 0.2])])
|
||||
|
||||
assert store.count(COLLECTION) == 0
|
||||
|
||||
|
||||
def test_manifest_dimensions_must_agree_with_the_collection_spec():
|
||||
with pytest.raises(ValueError, match="manifest declares 8 dimensions"):
|
||||
ChunkLoader(InMemoryVectorStore(), spec(), manifest(dimensions=8))
|
||||
|
||||
|
||||
def test_collection_spec_refuses_a_nonpositive_vector_size():
|
||||
with pytest.raises(ValueError, match="vector_size must be positive"):
|
||||
CollectionSpec(name=COLLECTION, vector_size=0)
|
||||
|
||||
|
||||
# --- A6: the corpus binding gate -----------------------------------------
|
||||
|
||||
|
||||
def test_the_manifest_lives_beside_the_data_so_the_point_count_stays_exact():
|
||||
store = InMemoryVectorStore()
|
||||
|
||||
loader(store).load(pairs(3))
|
||||
|
||||
assert store.count(COLLECTION) == 3, "the manifest must not inflate the count"
|
||||
assert store.count(manifest_collection(COLLECTION)) == 1
|
||||
stored = read_manifest(store, COLLECTION)
|
||||
assert stored.corpus_sha256 == CORPUS_SHA
|
||||
assert stored.model_id == MODEL
|
||||
assert stored.provider == "titan-v2"
|
||||
|
||||
|
||||
def test_a_second_corpus_generation_is_refused_and_nothing_is_written():
|
||||
store = InMemoryVectorStore()
|
||||
loader(store).load(pairs(3))
|
||||
|
||||
with pytest.raises(CorpusMismatch, match="does not match"):
|
||||
loader(store, manifest=manifest(corpus_sha256=OTHER_SHA)).load(pairs(2))
|
||||
|
||||
assert store.count(COLLECTION) == 3
|
||||
assert read_manifest(store, COLLECTION).corpus_sha256 == CORPUS_SHA
|
||||
|
||||
|
||||
def test_a_second_model_is_refused_even_when_the_corpus_matches():
|
||||
store = InMemoryVectorStore()
|
||||
loader(store).load(pairs(1))
|
||||
|
||||
with pytest.raises(CorpusMismatch, match="cohere.embed-v4:0"):
|
||||
loader(store, manifest=manifest(model_id="cohere.embed-v4:0")).load(pairs(1))
|
||||
|
||||
|
||||
def test_a_query_subspace_vector_is_refused_for_a_document_collection():
|
||||
store = InMemoryVectorStore()
|
||||
loader(store).load(pairs(1))
|
||||
|
||||
with pytest.raises(CorpusMismatch, match="input kind query"):
|
||||
loader(store, manifest=manifest(input_kind="query")).load(pairs(1))
|
||||
|
||||
|
||||
def test_a_dimension_change_is_refused():
|
||||
store = InMemoryVectorStore()
|
||||
loader(store).load(pairs(1))
|
||||
|
||||
eight = CollectionSpec(name=COLLECTION, vector_size=8)
|
||||
with pytest.raises(CorpusMismatch, match="8 dimensions"):
|
||||
ChunkLoader(store, eight, manifest(dimensions=8)).load(
|
||||
[(chunk_record(), [0.1] * 8)]
|
||||
)
|
||||
|
||||
|
||||
def test_an_existing_collection_with_no_manifest_is_refused():
|
||||
store = InMemoryVectorStore()
|
||||
store.create_collection(spec())
|
||||
store.upsert(COLLECTION, [build_point(chunk_record(), vector())])
|
||||
|
||||
with pytest.raises(CorpusMismatch, match="has no manifest"):
|
||||
loader(store).load(pairs(1))
|
||||
|
||||
|
||||
def test_the_same_corpus_and_model_is_allowed_through():
|
||||
store = InMemoryVectorStore()
|
||||
loader(store).load(pairs(3))
|
||||
|
||||
report = loader(store).load(pairs(3))
|
||||
|
||||
assert report.collection_count == 3
|
||||
assert report.count_matches is True
|
||||
|
||||
|
||||
# --- the v1 point-count gate ---------------------------------------------
|
||||
|
||||
|
||||
def test_assert_point_count_passes_when_every_chunk_has_exactly_one_point():
|
||||
store = InMemoryVectorStore()
|
||||
active = loader(store)
|
||||
active.load(pairs(3))
|
||||
|
||||
assert active.assert_point_count(3) == 3
|
||||
|
||||
|
||||
def test_assert_point_count_raises_when_the_collection_is_short():
|
||||
store = InMemoryVectorStore()
|
||||
active = loader(store)
|
||||
active.load(pairs(3))
|
||||
|
||||
with pytest.raises(PointCountMismatch, match="holds 3 points but the corpus has 4"):
|
||||
active.assert_point_count(4)
|
||||
|
||||
|
||||
# --- mode A: exhaustive filter retrieval ---------------------------------
|
||||
|
||||
|
||||
def section_pairs(drug_id, section_key, parts):
|
||||
return [
|
||||
(
|
||||
chunk_record(
|
||||
chunk_id=f"{drug_id}__{section_key}__{i}",
|
||||
drug_id=drug_id,
|
||||
section_key=section_key,
|
||||
part_index=i,
|
||||
part_count=parts,
|
||||
),
|
||||
vector(0.1),
|
||||
)
|
||||
for i in range(parts)
|
||||
]
|
||||
|
||||
|
||||
def test_a_filter_returns_every_part_of_a_section_not_a_top_k():
|
||||
"""The rule mode A exists for: two of five contraindications is worse than none."""
|
||||
store = InMemoryVectorStore()
|
||||
loader(store).load(
|
||||
section_pairs("metformin", "chong_chi_dinh", 5)
|
||||
+ section_pairs("metformin", "lieu_luong_va_cach_dung", 3)
|
||||
+ section_pairs("pantoprazol", "chong_chi_dinh", 2)
|
||||
)
|
||||
|
||||
found = store.find_by_payload(
|
||||
COLLECTION, {"drug_id": "metformin", "section_key": "chong_chi_dinh"}
|
||||
)
|
||||
|
||||
assert len(found) == 5
|
||||
assert sorted(p.payload["part_index"] for p in found) == [0, 1, 2, 3, 4]
|
||||
assert {p.payload["drug_id"] for p in found} == {"metformin"}
|
||||
|
||||
|
||||
def test_a_filter_never_leaks_a_neighbouring_drugs_section():
|
||||
store = InMemoryVectorStore()
|
||||
loader(store).load(
|
||||
section_pairs("pantoprazol", "chong_chi_dinh", 2)
|
||||
+ section_pairs("omeprazol", "chong_chi_dinh", 2)
|
||||
)
|
||||
|
||||
found = store.find_by_payload(
|
||||
COLLECTION, {"drug_id": "pantoprazol", "section_key": "chong_chi_dinh"}
|
||||
)
|
||||
|
||||
assert len(found) == 2
|
||||
assert {p.payload["drug_id"] for p in found} == {"pantoprazol"}
|
||||
|
||||
|
||||
def test_a_list_valued_field_matches_on_any_element():
|
||||
store = InMemoryVectorStore()
|
||||
record = chunk_record(atc_codes=["A10BA02", "A10BD20"])
|
||||
loader(store).load([(record, vector())])
|
||||
|
||||
assert len(store.find_by_payload(COLLECTION, {"atc_codes": "A10BD20"})) == 1
|
||||
assert len(store.find_by_payload(COLLECTION, {"atc_codes": "J05AF06"})) == 0
|
||||
|
||||
|
||||
def test_a_filter_matching_nothing_returns_empty_rather_than_raising():
|
||||
store = InMemoryVectorStore()
|
||||
loader(store).load(pairs(2))
|
||||
|
||||
assert store.find_by_payload(COLLECTION, {"drug_id": "khong_ton_tai"}) == []
|
||||
|
||||
|
||||
def test_a_filter_with_no_condition_is_refused():
|
||||
store = InMemoryVectorStore()
|
||||
loader(store).load(pairs(1))
|
||||
|
||||
with pytest.raises(ValueError, match="at least one condition"):
|
||||
store.find_by_payload(COLLECTION, {})
|
||||
|
||||
|
||||
def test_parts_reassemble_in_order_into_the_whole_section():
|
||||
store = InMemoryVectorStore()
|
||||
bodies = ["Phần một.", "Phần hai.", "Phần ba."]
|
||||
records = [
|
||||
(
|
||||
chunk_record(
|
||||
chunk_id=f"metformin__chong_chi_dinh__{i}",
|
||||
drug_id="metformin",
|
||||
section_key="chong_chi_dinh",
|
||||
text=body,
|
||||
part_index=i,
|
||||
part_count=len(bodies),
|
||||
),
|
||||
vector(),
|
||||
)
|
||||
for i, body in enumerate(bodies)
|
||||
]
|
||||
loader(store).load(records)
|
||||
|
||||
found = store.find_by_payload(
|
||||
COLLECTION, {"drug_id": "metformin", "section_key": "chong_chi_dinh"}
|
||||
)
|
||||
ordered = sorted(found, key=lambda p: p.payload["part_index"])
|
||||
|
||||
assert [p.payload["text"] for p in ordered] == bodies
|
||||
assert {p.payload["part_count"] for p in found} == {3}
|
||||
|
||||
|
||||
# --- corpus digest and reading -------------------------------------------
|
||||
|
||||
|
||||
def test_corpus_sha256_changes_when_a_single_byte_changes(tmp_path):
|
||||
path = tmp_path / "chunks.jsonl"
|
||||
path.write_text(json.dumps(chunk_record()) + "\n", encoding="utf-8")
|
||||
before = corpus_sha256(path)
|
||||
|
||||
path.write_text(
|
||||
json.dumps(chunk_record(text="Người lớn: 301 mg.")) + "\n", encoding="utf-8"
|
||||
)
|
||||
|
||||
assert corpus_sha256(path) != before
|
||||
assert len(before) == 64
|
||||
|
||||
|
||||
def test_the_same_data_hashes_the_same_under_crlf_and_lf(tmp_path):
|
||||
"""A gate that cries wolf gets switched off.
|
||||
|
||||
A raw-byte digest made a Windows CRLF checkout and a Linux LF checkout of
|
||||
identical data disagree, so A6 would refuse a CI load against the very
|
||||
corpus it was built from.
|
||||
"""
|
||||
body = json.dumps(chunk_record()) + "\n" + json.dumps(chunk_record("b")) + "\n"
|
||||
lf = tmp_path / "lf.jsonl"
|
||||
crlf = tmp_path / "crlf.jsonl"
|
||||
lf.write_bytes(body.encode("utf-8"))
|
||||
crlf.write_bytes(body.replace("\n", "\r\n").encode("utf-8"))
|
||||
|
||||
assert corpus_sha256(lf) == corpus_sha256(crlf)
|
||||
assert crlf.stat().st_size > lf.stat().st_size, "the files really do differ"
|
||||
|
||||
|
||||
def test_blank_lines_are_skipped_and_bad_json_names_its_line(tmp_path):
|
||||
path = tmp_path / "chunks.jsonl"
|
||||
path.write_text(
|
||||
json.dumps(chunk_record()) + "\n\n" + json.dumps(chunk_record("b")) + "\n",
|
||||
encoding="utf-8",
|
||||
)
|
||||
|
||||
assert count_chunks(path) == 2
|
||||
|
||||
broken = tmp_path / "broken.jsonl"
|
||||
broken.write_text(json.dumps(chunk_record()) + "\n{oops\n", encoding="utf-8")
|
||||
with pytest.raises(ValueError, match="line 2 is not valid JSON"):
|
||||
list(iter_chunk_records(broken))
|
||||
|
||||
|
||||
# --- the real artifact ----------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not REAL_CHUNKS.exists(), reason="chunks.jsonl has not been generated"
|
||||
)
|
||||
def test_every_real_chunk_satisfies_the_loader_provenance_contract():
|
||||
"""Whole-artifact scope: all records in `data/processed/chunks.jsonl`."""
|
||||
seen_ids = set()
|
||||
seen_points = set()
|
||||
total = 0
|
||||
for record in iter_chunk_records(REAL_CHUNKS):
|
||||
validate_chunk_record(record)
|
||||
point = point_id_for(record["chunk_id"])
|
||||
assert point not in seen_points, f"point id collision on {record['chunk_id']}"
|
||||
seen_points.add(point)
|
||||
seen_ids.add(record["chunk_id"])
|
||||
total += 1
|
||||
|
||||
assert total == len(seen_ids), "duplicate chunk_id in the artifact"
|
||||
assert total == len(seen_points)
|
||||
# A floor, not the exact count: the corpus is regenerated as `segment/`
|
||||
# changes, but a truncated or half-written artifact must not pass as
|
||||
# whole-artifact evidence. Measured 15,066 records on 2026-08-04.
|
||||
assert total > 10_000, f"chunks.jsonl looks truncated: only {total} records"
|
||||
@@ -0,0 +1,377 @@
|
||||
"""`QdrantVectorStore` against a real Qdrant, skipped when none is running.
|
||||
|
||||
The rest of the load suite runs against `InMemoryVectorStore` and proves the
|
||||
loader's rules. It cannot prove the adapter: whether Qdrant accepts a uuid5
|
||||
string as a point id, whether `create_payload_index` takes a bare `"keyword"`,
|
||||
whether an upsert of an existing id replaces rather than appends. Those are
|
||||
claims about another system, and the same class of claim as the Bedrock request
|
||||
bodies that are still documentation-derived and unproven — so they get a live
|
||||
check, against a free local container rather than a paid API.
|
||||
|
||||
Start one with `docker compose -f infra/docker/docker-compose.yml up -d qdrant`.
|
||||
Without it these tests skip; they never fail for being offline.
|
||||
|
||||
Every test works in its own collection named after the test and deletes it
|
||||
afterwards, so a shared local Qdrant is not left holding fixtures.
|
||||
"""
|
||||
import os
|
||||
import uuid
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from ingestion.load import (
|
||||
ChunkLoader,
|
||||
CollectionSpec,
|
||||
CorpusManifest,
|
||||
CorpusMismatch,
|
||||
manifest_collection,
|
||||
point_id_for,
|
||||
read_manifest,
|
||||
)
|
||||
from ingestion.load.qdrant_repo import DEFAULT_URL, SCROLL_PAGE, QdrantVectorStore
|
||||
|
||||
DIMENSIONS = 4
|
||||
QDRANT_URL = os.environ.get("QDRANT_URL", DEFAULT_URL)
|
||||
REAL_CHUNKS = (
|
||||
Path(__file__).resolve().parents[1] / "data" / "processed" / "chunks.jsonl"
|
||||
)
|
||||
|
||||
pytestmark = pytest.mark.integration
|
||||
|
||||
|
||||
def _server_is_up() -> bool:
|
||||
try:
|
||||
from qdrant_client import QdrantClient
|
||||
except ImportError:
|
||||
return False
|
||||
try:
|
||||
QdrantClient(url=QDRANT_URL, timeout=3.0).get_collections()
|
||||
except Exception:
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
requires_qdrant = pytest.mark.skipif(
|
||||
not _server_is_up(), reason=f"no Qdrant reachable at {QDRANT_URL}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def store():
|
||||
return QdrantVectorStore(url=QDRANT_URL, timeout=10.0)
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def collection(store):
|
||||
name = f"test_load_{uuid.uuid4().hex[:10]}"
|
||||
yield name
|
||||
for target in (manifest_collection(name), name):
|
||||
if store.collection_exists(target):
|
||||
store.delete_collection(target)
|
||||
|
||||
|
||||
def chunk_record(chunk_id, **overrides):
|
||||
record = {
|
||||
"schema_version": 4,
|
||||
"chunk_id": chunk_id,
|
||||
"drug_id": "abacavir",
|
||||
"drug_name": "ABACAVIR",
|
||||
"section_key": "lieu_luong_va_cach_dung",
|
||||
"section_display_name": "Liều lượng và cách dùng",
|
||||
"text": "Người lớn: 300 mg, hai lần mỗi ngày.",
|
||||
"source_text": "Người lớn: 300 mg, hai lần mỗi ngày.",
|
||||
"heading_physical_page": 100,
|
||||
"source_page_range": [100, 102],
|
||||
"printed_page_range": [101, 103],
|
||||
"atc_codes": ["J05AF06"],
|
||||
"part_index": 0,
|
||||
"part_count": 1,
|
||||
"est_tokens": 14,
|
||||
"oversized": False,
|
||||
"chunk_kind": "prose",
|
||||
"attachments": [],
|
||||
"has_quarantined_content": False,
|
||||
}
|
||||
record.update(overrides)
|
||||
return record
|
||||
|
||||
|
||||
def manifest(**overrides):
|
||||
values = {
|
||||
"corpus_sha256": "a" * 64,
|
||||
"chunk_count": 2,
|
||||
"model_id": "amazon.titan-embed-text-v2:0",
|
||||
"dimensions": DIMENSIONS,
|
||||
"input_kind": "document",
|
||||
"provider": "titan-v2",
|
||||
}
|
||||
values.update(overrides)
|
||||
return CorpusManifest(**values)
|
||||
|
||||
|
||||
def pairs(count):
|
||||
return [
|
||||
(chunk_record(f"abacavir__section__{i}"), [0.1 * (i + 1)] * DIMENSIONS)
|
||||
for i in range(count)
|
||||
]
|
||||
|
||||
|
||||
@requires_qdrant
|
||||
def test_a_real_load_creates_the_collection_indexes_and_points(store, collection):
|
||||
spec = CollectionSpec(name=collection, vector_size=DIMENSIONS)
|
||||
|
||||
report = ChunkLoader(store, spec, manifest()).load(pairs(2))
|
||||
|
||||
assert report.collection_created is True
|
||||
assert report.collection_count == 2
|
||||
assert store.collection_exists(collection)
|
||||
assert store.count(collection) == 2
|
||||
|
||||
|
||||
@requires_qdrant
|
||||
def test_qdrant_accepts_the_derived_uuid5_point_id_and_returns_the_payload(
|
||||
store, collection
|
||||
):
|
||||
spec = CollectionSpec(name=collection, vector_size=DIMENSIONS)
|
||||
record = chunk_record("abacavir__lieu_luong__0")
|
||||
ChunkLoader(store, spec, manifest()).load([(record, [0.5] * DIMENSIONS)])
|
||||
|
||||
stored = store.retrieve(collection, point_id_for(record["chunk_id"]))
|
||||
|
||||
assert stored is not None
|
||||
assert stored.id == point_id_for(record["chunk_id"])
|
||||
assert stored.payload["chunk_id"] == record["chunk_id"]
|
||||
assert stored.payload["source_page_range"] == [100, 102]
|
||||
assert stored.payload["atc_codes"] == ["J05AF06"]
|
||||
assert stored.payload["has_quarantined_content"] is False
|
||||
assert len(stored.vector) == DIMENSIONS
|
||||
|
||||
|
||||
@requires_qdrant
|
||||
def test_loading_twice_against_a_real_server_does_not_duplicate(store, collection):
|
||||
spec = CollectionSpec(name=collection, vector_size=DIMENSIONS)
|
||||
data = pairs(2)
|
||||
|
||||
ChunkLoader(store, spec, manifest()).load(data)
|
||||
second = ChunkLoader(store, spec, manifest()).load(data)
|
||||
|
||||
assert second.collection_created is False
|
||||
assert store.count(collection) == 2, "a re-run duplicated points in Qdrant"
|
||||
|
||||
|
||||
@requires_qdrant
|
||||
def test_the_manifest_sidecar_round_trips_and_leaves_the_count_exact(
|
||||
store, collection
|
||||
):
|
||||
spec = CollectionSpec(name=collection, vector_size=DIMENSIONS)
|
||||
ChunkLoader(store, spec, manifest()).load(pairs(2))
|
||||
|
||||
stored = read_manifest(store, collection)
|
||||
|
||||
assert stored is not None
|
||||
assert stored.corpus_sha256 == "a" * 64
|
||||
assert stored.model_id == "amazon.titan-embed-text-v2:0"
|
||||
assert stored.dimensions == DIMENSIONS
|
||||
assert store.count(collection) == 2
|
||||
assert store.count(manifest_collection(collection)) == 1
|
||||
|
||||
|
||||
@requires_qdrant
|
||||
def test_a_second_corpus_is_refused_against_a_real_collection(store, collection):
|
||||
spec = CollectionSpec(name=collection, vector_size=DIMENSIONS)
|
||||
ChunkLoader(store, spec, manifest()).load(pairs(2))
|
||||
|
||||
with pytest.raises(CorpusMismatch, match="does not match"):
|
||||
ChunkLoader(store, spec, manifest(corpus_sha256="b" * 64)).load(pairs(2))
|
||||
|
||||
assert store.count(collection) == 2
|
||||
|
||||
|
||||
@requires_qdrant
|
||||
def test_retrieve_returns_none_for_an_id_that_was_never_loaded(store, collection):
|
||||
spec = CollectionSpec(name=collection, vector_size=DIMENSIONS)
|
||||
ChunkLoader(store, spec, manifest()).load(pairs(1))
|
||||
|
||||
assert store.retrieve(collection, point_id_for("never__loaded__0")) is None
|
||||
|
||||
|
||||
@requires_qdrant
|
||||
def test_bbox_floats_lose_precision_in_qdrant_but_nothing_else_does(store, collection):
|
||||
"""Pins a measured round-trip loss so it cannot silently get worse.
|
||||
|
||||
Scrolling all 15,066 points of a full load on 2026-08-04 found 86 chunks
|
||||
whose payload did not compare equal to its source record. Every one of the
|
||||
96 differing leaf values was a float inside `attachments[].bbox`, the
|
||||
largest delta was 5.684e-14, and **no** text, id, page number, page range,
|
||||
token count or boolean differed at all. A PDF point is 1/72 inch, so that
|
||||
delta cannot move a rendered crop; what would matter is the loss spreading
|
||||
to another field, or growing. This test fails if either happens.
|
||||
"""
|
||||
spec = CollectionSpec(name=collection, vector_size=DIMENSIONS)
|
||||
# 17 significant digits: the real corpus carries these, and they are what
|
||||
# does not survive a float64 -> JSON -> float64 round trip.
|
||||
bbox = [44.45098876953125, 397.45245361328125, 278.09100341796875, 463.4044494628906]
|
||||
record = chunk_record(
|
||||
"cefazolin__lieu_luong_va_cach_dung__0",
|
||||
has_quarantined_content=True,
|
||||
attachments=[
|
||||
{
|
||||
"block_id": "p344_t2",
|
||||
"kind": "table",
|
||||
"shape": "simple_table",
|
||||
"physical_page": 344,
|
||||
"bbox": bbox,
|
||||
"quarantined": True,
|
||||
"header_row": ["Cỡ lọ", "Lượng\ndung môi"],
|
||||
}
|
||||
],
|
||||
)
|
||||
ChunkLoader(store, spec, manifest()).load([(record, [0.3] * DIMENSIONS)])
|
||||
|
||||
stored = store.retrieve(collection, point_id_for(record["chunk_id"])).payload
|
||||
attachment = stored["attachments"][0]
|
||||
|
||||
for value, original in zip(attachment["bbox"], bbox, strict=True):
|
||||
assert abs(value - original) < 1e-9, "bbox drift grew beyond rounding"
|
||||
|
||||
assert attachment["block_id"] == "p344_t2"
|
||||
assert attachment["physical_page"] == 344
|
||||
assert attachment["quarantined"] is True
|
||||
assert attachment["header_row"] == ["Cỡ lọ", "Lượng\ndung môi"]
|
||||
for field in (
|
||||
"chunk_id",
|
||||
"drug_id",
|
||||
"drug_name",
|
||||
"section_key",
|
||||
"text",
|
||||
"heading_physical_page",
|
||||
"source_page_range",
|
||||
"atc_codes",
|
||||
"est_tokens",
|
||||
"chunk_kind",
|
||||
"has_quarantined_content",
|
||||
):
|
||||
assert stored[field] == record[field], f"{field} must round-trip exactly"
|
||||
|
||||
|
||||
@requires_qdrant
|
||||
def test_the_payload_index_actually_serves_a_filtered_query(store, collection):
|
||||
"""Creating an index proves nothing; querying through it does.
|
||||
|
||||
Mode A of the delivery plan never ranks by vector — it filters on
|
||||
`drug_id` + `section_key` and returns the whole section. Until this test
|
||||
existed the loader had only established that `create_payload_index`
|
||||
returned without error.
|
||||
"""
|
||||
spec = CollectionSpec(name=collection, vector_size=DIMENSIONS)
|
||||
data = (
|
||||
_section("metformin", "chong_chi_dinh", 5)
|
||||
+ _section("metformin", "lieu_luong_va_cach_dung", 3)
|
||||
+ _section("pantoprazol", "chong_chi_dinh", 4)
|
||||
)
|
||||
ChunkLoader(store, spec, manifest()).load(data)
|
||||
|
||||
found = store.find_by_payload(
|
||||
collection, {"drug_id": "metformin", "section_key": "chong_chi_dinh"}
|
||||
)
|
||||
|
||||
assert len(found) == 5
|
||||
assert sorted(p.payload["part_index"] for p in found) == [0, 1, 2, 3, 4]
|
||||
assert {p.payload["drug_id"] for p in found} == {"metformin"}
|
||||
|
||||
|
||||
@requires_qdrant
|
||||
def test_a_section_longer_than_one_scroll_page_comes_back_whole(store, collection):
|
||||
"""Paging must not truncate a section — that is the mode A failure mode.
|
||||
|
||||
Sized deliberately above `SCROLL_PAGE` (256) so a single-page implementation
|
||||
fails here rather than in production on the one drug with a long section.
|
||||
"""
|
||||
spec = CollectionSpec(name=collection, vector_size=DIMENSIONS)
|
||||
parts = SCROLL_PAGE + 44
|
||||
ChunkLoader(store, spec, manifest()).load(
|
||||
_section("insulin", "lieu_luong_va_cach_dung", parts)
|
||||
)
|
||||
|
||||
found = store.find_by_payload(
|
||||
collection,
|
||||
{"drug_id": "insulin", "section_key": "lieu_luong_va_cach_dung"},
|
||||
)
|
||||
|
||||
assert len(found) == parts
|
||||
assert sorted(p.payload["part_index"] for p in found) == list(range(parts))
|
||||
|
||||
|
||||
@requires_qdrant
|
||||
def test_a_list_valued_atc_field_matches_on_any_element_in_qdrant(store, collection):
|
||||
spec = CollectionSpec(name=collection, vector_size=DIMENSIONS)
|
||||
record = chunk_record("metformin__lieu_luong__0", atc_codes=["A10BA02", "A10BD20"])
|
||||
ChunkLoader(store, spec, manifest()).load([(record, [0.4] * DIMENSIONS)])
|
||||
|
||||
assert len(store.find_by_payload(collection, {"atc_codes": "A10BD20"})) == 1
|
||||
assert len(store.find_by_payload(collection, {"atc_codes": "J05AF06"})) == 0
|
||||
|
||||
|
||||
@requires_qdrant
|
||||
@pytest.mark.skipif(
|
||||
not REAL_CHUNKS.exists(), reason="chunks.jsonl has not been generated"
|
||||
)
|
||||
def test_a_real_multipart_section_round_trips_through_the_filter(store, collection):
|
||||
"""Against the real artifact, not fixtures: every part, and only those."""
|
||||
from ingestion.load import iter_chunk_records
|
||||
|
||||
wanted = None
|
||||
records = []
|
||||
for record in iter_chunk_records(REAL_CHUNKS):
|
||||
if wanted is None and record["part_count"] >= 4:
|
||||
wanted = (record["drug_id"], record["section_key"])
|
||||
records.append(record)
|
||||
assert wanted is not None, "no multi-part section in the artifact"
|
||||
|
||||
drug_id, section_key = wanted
|
||||
expected = {
|
||||
r["chunk_id"]
|
||||
for r in records
|
||||
if r["drug_id"] == drug_id and r["section_key"] == section_key
|
||||
}
|
||||
subset = [
|
||||
r for r in records
|
||||
if r["drug_id"] == drug_id or r["chunk_id"].startswith("abacavir")
|
||||
]
|
||||
|
||||
spec = CollectionSpec(name=collection, vector_size=DIMENSIONS)
|
||||
ChunkLoader(store, spec, manifest()).load(
|
||||
(r, [0.2] * DIMENSIONS) for r in subset
|
||||
)
|
||||
|
||||
found = store.find_by_payload(
|
||||
collection, {"drug_id": drug_id, "section_key": section_key}
|
||||
)
|
||||
|
||||
assert {p.payload["chunk_id"] for p in found} == expected
|
||||
assert len(expected) >= 4
|
||||
|
||||
|
||||
def _section(drug_id, section_key, parts):
|
||||
return [
|
||||
(
|
||||
chunk_record(
|
||||
f"{drug_id}__{section_key}__{i}",
|
||||
drug_id=drug_id,
|
||||
section_key=section_key,
|
||||
part_index=i,
|
||||
part_count=parts,
|
||||
),
|
||||
[0.1] * DIMENSIONS,
|
||||
)
|
||||
for i in range(parts)
|
||||
]
|
||||
|
||||
|
||||
@requires_qdrant
|
||||
def test_an_unsupported_distance_is_rejected_before_the_server_is_called(store):
|
||||
with pytest.raises(ValueError, match="unsupported distance"):
|
||||
store.create_collection(
|
||||
CollectionSpec(name="never_created", vector_size=4, distance="manhattan")
|
||||
)
|
||||
@@ -1,3 +1,5 @@
|
||||
from dataclasses import replace
|
||||
|
||||
import pytest
|
||||
|
||||
from ingestion.extract.models import Span
|
||||
@@ -50,6 +52,55 @@ def test_non_bold_combined_heading_value_span_confirmed_real_amitriptylin_case()
|
||||
m = list(assemble(spans))[0]
|
||||
assert m.sections["ma_atc"].text == "N06AA09."
|
||||
assert m.atc_codes == ["N06AA09"]
|
||||
part = m.sections["ma_atc"].parts[0]
|
||||
assert part.physical_page == 184
|
||||
assert part.bbox != [0.0, 0.0, 0.0, 0.0]
|
||||
assert part.source_span_ids == [spans[3].span_id]
|
||||
|
||||
|
||||
def test_combined_international_name_and_atc_heading_is_a_title_anchor():
|
||||
# Confirmed real GnRH class-monograph variant on physical page 1371.
|
||||
spans = [
|
||||
_span("THUỐC TƯƠNG TỰ HORMON GIẢI PHÓNG", 1371, 664.0),
|
||||
_span("GONADOTROPIN", 1371, 676.0),
|
||||
_span("Tên chung quốc tế và mã ATC", 1371, 690.0),
|
||||
_span("Gonadorelin: H01CA01; Triptorelin: L02AE04.", 1371, 702.0, bold=False),
|
||||
_span("Chỉ định", 1371, 714.0),
|
||||
_span("Kích thích phóng noãn.", 1371, 726.0, bold=False),
|
||||
]
|
||||
m = list(assemble(spans))[0]
|
||||
assert m.drug_name == "THUỐC TƯƠNG TỰ HORMON GIẢI PHÓNG GONADOTROPIN"
|
||||
assert m.atc_codes == ["H01CA01", "L02AE04"]
|
||||
|
||||
|
||||
def test_plain_wrapped_section_label_is_body_not_a_heading():
|
||||
# NADROPARIN CALCI p1016: "không phải là" / "chống chỉ định."
|
||||
# are adjacent lines of one sentence in the same PDF block.
|
||||
spans = [
|
||||
_span("NADROPARIN CALCI", 1016, 60.0),
|
||||
_span("Tên chung quốc tế", 1016, 80.0),
|
||||
_span("Nadroparin calcium.", 1016, 92.0, bold=False),
|
||||
_span("Thời kỳ cho con bú", 1016, 110.0),
|
||||
]
|
||||
lead = replace(_span("Việc dùng thuốc không phải là", 1016, 122.0, bold=False),
|
||||
block=4, line=7)
|
||||
tail = replace(_span("chống chỉ định.", 1016, 134.0, bold=False),
|
||||
block=4, line=8)
|
||||
m = list(assemble(spans + [lead, tail]))[0]
|
||||
assert m.sections["thoi_ky_cho_con_bu"].text.endswith("chống chỉ định.")
|
||||
assert "chong_chi_dinh" not in m.sections
|
||||
|
||||
|
||||
def test_plain_heading_after_completed_prose_still_opens_section():
|
||||
spans = [
|
||||
_span("TESTDRUG", 300, 60.0),
|
||||
_span("Tên chung quốc tế", 300, 80.0),
|
||||
replace(_span("Testdrug.", 300, 92.0, bold=False), block=2, line=0),
|
||||
replace(_span("Chỉ định", 300, 104.0, bold=False), block=2, line=1),
|
||||
replace(_span("Điều trị thử nghiệm.", 300, 116.0, bold=False), block=2, line=2),
|
||||
]
|
||||
m = list(assemble(spans))[0]
|
||||
assert m.sections["chi_dinh"].text == "Điều trị thử nghiệm."
|
||||
|
||||
|
||||
def test_atc_stated_absent_propagates():
|
||||
@@ -330,3 +381,71 @@ def test_a_bold_label_line_still_opens_its_section():
|
||||
]
|
||||
monograph = list(assemble(spans))[0]
|
||||
assert monograph.sections["chong_chi_dinh"].text == "Suy tủy nặng."
|
||||
|
||||
|
||||
def test_a_section_name_printed_mid_line_is_body_not_a_heading():
|
||||
"""CISPLATIN, physical page 402 — confirmed content loss.
|
||||
|
||||
The book prints "Suy thận: Chống chỉ định." inside the dosing section. The
|
||||
second half is itself a section name, so it was matched as a heading: the
|
||||
renal-impairment contraindication vanished from the dosing text and the
|
||||
section ended on a bare "Suy thận:". ISOPRENALIN had the same shape. A
|
||||
real heading opens its line; this one does not.
|
||||
"""
|
||||
spans = [
|
||||
_span("CISPLATIN", 401, 60.0),
|
||||
_span("Tên chung quốc tế", 401, 80.0),
|
||||
_span("Cisplatinum.", 401, 92.0, bold=False),
|
||||
_span("Liều lượng và cách dùng", 401, 110.0),
|
||||
_span("Truyền tĩnh mạch mỗi 3 tuần.", 401, 122.0, bold=False),
|
||||
]
|
||||
label = _span("Suy thận: ", 401, 140.0, bold=False)
|
||||
label = replace(label, block=4, line=0, x0=35.0, x1=70.0)
|
||||
trailing = _span("Chống chỉ định.", 401, 140.0, bold=False)
|
||||
trailing = replace(trailing, block=4, line=0, x0=70.0, x1=140.0)
|
||||
|
||||
monograph = list(assemble(spans + [label, trailing]))[0]
|
||||
dosing = monograph.sections["lieu_luong_va_cach_dung"].text
|
||||
assert "Suy thận: Chống chỉ định." in dosing
|
||||
assert "chong_chi_dinh" not in monograph.sections
|
||||
|
||||
|
||||
def test_italic_cross_reference_overlapping_its_neighbour_by_a_hairline_is_body():
|
||||
"""NEVIRAPIN, physical page 1045 — confirmed misassignment, whole-corpus.
|
||||
|
||||
The book prints `Xem thêm mục ` (x1=104.89) immediately before an italic
|
||||
`Liều lượng và cách dùng` (x0=104.88): the trailing space's advance width
|
||||
makes the neighbour end 0.01pt *after* the cross-reference starts. An
|
||||
end-before-start test therefore read a mid-line cross-reference as a
|
||||
heading. Same shape, same cause, in CALCI LACTAT (p296, `xem thêm mục
|
||||
Tương tác thuốc`, 0.02pt) and CEFAZOLIN (p344, `ghi ở mục: Dạng thuốc và
|
||||
hàm lượng.`), where 4,533 characters of adult dosing were filed under
|
||||
dosage forms.
|
||||
"""
|
||||
spans = [
|
||||
_span("NEVIRAPIN", 1044, 60.0),
|
||||
_span("Tên chung quốc tế", 1044, 80.0),
|
||||
_span("Nevirapine.", 1044, 92.0, bold=False),
|
||||
_span("Hướng dẫn cách xử trí ADR", 1044, 110.0),
|
||||
_span("Điều trị các phản ứng bất lợi theo triệu chứng.", 1044, 122.0, bold=False),
|
||||
]
|
||||
lead = replace(_span("Xem thêm mục ", 1044, 140.0, bold=False),
|
||||
block=4, line=0, x0=43.94, x1=104.89)
|
||||
reference = replace(_span("Liều lượng và cách dùng", 1044, 140.0, bold=False),
|
||||
block=4, line=0, x0=104.88, x1=199.57)
|
||||
|
||||
monograph = list(assemble(spans + [lead, reference]))[0]
|
||||
assert "Xem thêm mục Liều lượng và cách dùng" in monograph.sections["huong_dan_xu_tri_adr"].text
|
||||
assert "lieu_luong_va_cach_dung" not in monograph.sections
|
||||
|
||||
|
||||
def test_a_section_name_opening_its_own_line_is_still_a_heading():
|
||||
spans = [
|
||||
_span("CISPLATIN", 401, 60.0),
|
||||
_span("Tên chung quốc tế", 401, 80.0),
|
||||
_span("Cisplatinum.", 401, 92.0, bold=False),
|
||||
_span("Chống chỉ định", 401, 110.0),
|
||||
_span("Suy tủy nặng.", 401, 122.0, bold=False),
|
||||
]
|
||||
monograph = list(assemble(spans))[0]
|
||||
assert monograph.sections["chong_chi_dinh"].text == "Suy tủy nặng."
|
||||
|
||||
@@ -1,5 +1,9 @@
|
||||
from ingestion.extract.models import Span
|
||||
from ingestion.segment.detector import detect_monograph_titles, detect_section_headings
|
||||
from ingestion.segment.detector import (
|
||||
detect_monograph_titles,
|
||||
detect_section_headings,
|
||||
in_monograph_range,
|
||||
)
|
||||
|
||||
|
||||
def _span(text, physical_page, printed_page, y0=100.0, bold=True, size=10.0):
|
||||
@@ -86,3 +90,12 @@ def test_unknown_bold_text_not_matched_as_section():
|
||||
def test_section_heading_outside_monograph_range_excluded():
|
||||
spans = [_span("Chỉ định", 5, 6, bold=True, size=9.5)]
|
||||
assert list(detect_section_headings(spans)) == []
|
||||
|
||||
|
||||
def test_back_index_cannot_reenter_range_via_bad_inferred_printed_page():
|
||||
# Confirmed real failure: physical page 1655 of the back index was mapped
|
||||
# to printed page 1496, making its "Tương tác thuốc" entry extend
|
||||
# ZOLPIDEM's source range from page 1494 through page 1655.
|
||||
index_span = _span("Tương tác thuốc", 1655, 1496)
|
||||
assert in_monograph_range(index_span) is False
|
||||
assert list(detect_section_headings([index_span])) == []
|
||||
|
||||
@@ -1,6 +1,12 @@
|
||||
from ingestion.extract.models import Span
|
||||
from ingestion.segment import assemble
|
||||
from ingestion.tables import SHAPE_GRID_2D, SHAPE_SIMPLE, TableRegion, index_by_page
|
||||
from ingestion.tables import (
|
||||
SHAPE_FORMULA_2D,
|
||||
SHAPE_GRID_2D,
|
||||
SHAPE_SIMPLE,
|
||||
TableRegion,
|
||||
index_by_page,
|
||||
)
|
||||
|
||||
|
||||
def _span(text, page, y0, *, bold=False, x0=50.0, block=0, line=0, column="left"):
|
||||
@@ -54,6 +60,40 @@ def test_table_spans_are_lifted_out_of_section_prose():
|
||||
assert block.quarantined is True
|
||||
|
||||
|
||||
def test_section_named_table_cell_does_not_change_owning_section():
|
||||
# Confirmed in WARFARIN p1485 and IOBITRIDOL p826: a table column named
|
||||
# "Chỉ định" belongs to the dosing table; it is not a document
|
||||
# section heading and must not move the block into chi_dinh.
|
||||
spans = [
|
||||
_span("WARFARIN", 1485, 60.0, bold=True),
|
||||
_span("Tên chung quốc tế", 1485, 80.0, bold=True),
|
||||
_span("Warfarinum.", 1485, 92.0),
|
||||
_span("Liều lượng và cách dùng", 1485, 200.0, bold=True),
|
||||
_span("Chỉ định", 1485, 400.0, bold=True, block=5),
|
||||
_span("INR 2,0 - 3,0", 1485, 412.0, block=5),
|
||||
]
|
||||
region = TableRegion("p1485_t0", 1485, (40.0, 380.0, 400.0, 460.0), 2, 2, SHAPE_SIMPLE)
|
||||
m = list(assemble(spans, table_index=index_by_page([region])))[0]
|
||||
assert "chi_dinh" not in m.sections
|
||||
assert len(m.tables) == 1
|
||||
assert m.tables[0].section_key == "lieu_luong_va_cach_dung"
|
||||
|
||||
|
||||
def test_wide_formula_band_does_not_swallow_the_opposite_column():
|
||||
spans = _monograph_spans([
|
||||
_span("Công thức:", 109, 360.0),
|
||||
_span("Cl", 109, 400.0, x0=280.0, column="left", block=5),
|
||||
_span("Xem thêm Liều lượng và cách dùng", 109, 400.0,
|
||||
x0=310.0, column="right", block=6),
|
||||
])
|
||||
# Deliberately extends across the gutter, as verified formula bands do.
|
||||
region = TableRegion("p109_f0", 109, (40.0, 380.0, 390.0, 430.0),
|
||||
2, 1, SHAPE_FORMULA_2D)
|
||||
m = list(assemble(spans, table_index=index_by_page([region])))[0]
|
||||
assert m.tables[0].text == "Cl"
|
||||
assert "Xem thêm Liều lượng và cách dùng" in m.sections["dang_thuoc_va_ham_luong"].text
|
||||
|
||||
|
||||
def test_without_a_region_map_behaviour_is_unchanged():
|
||||
spans = _monograph_spans([
|
||||
_span("Thuốc dùng đường uống.", 109, 220.0),
|
||||
@@ -85,24 +125,53 @@ def test_non_table_regions_are_never_lifted():
|
||||
|
||||
|
||||
def test_table_block_ids_stay_unique_when_a_section_resumes():
|
||||
# a region flushed twice (section closes, then resumes) must not emit two
|
||||
# blocks with the same table_id — provenance ids have to be unique
|
||||
# Confirmed on CAPECITABIN pp. 308-309 and IMATINIB p. 795: PDF block
|
||||
# order can place a visually later heading between cells from one physical
|
||||
# table. The complete region must stay atomic and owned by the section
|
||||
# active where the table first appears.
|
||||
spans = [
|
||||
_span("CEFAMANDOL", 339, 60.0, bold=True),
|
||||
_span("Tên chung quốc tế", 339, 80.0, bold=True),
|
||||
_span("Cefamandolum.", 339, 92.0),
|
||||
_span("Liều lượng và cách dùng", 339, 200.0, bold=True),
|
||||
_span("80 - 50", 339, 400.0, block=5),
|
||||
_span("Liều lượng và cách dùng", 339, 500.0, bold=True),
|
||||
# Visually below the table, but emitted before its final cell by the
|
||||
# PDF's internal block order.
|
||||
_span("Tương tác thuốc", 339, 640.0, bold=True),
|
||||
_span("< 25 - 10", 339, 600.0, block=9),
|
||||
_span("Không phối hợp với thuốc X.", 339, 660.0, block=10),
|
||||
]
|
||||
region = TableRegion("p339_t0", 339, (40.0, 380.0, 400.0, 620.0), 5, 2, SHAPE_SIMPLE)
|
||||
m = list(assemble(spans, table_index=index_by_page([region])))[0]
|
||||
# table_id is deterministic per REGION, so two parts of one table share
|
||||
# it on purpose; table_part_id is the unique key, derived from the first
|
||||
# source span rather than a counter (a counter would renumber whenever
|
||||
# anything upstream shifted, hiding rather than identifying a duplicate)
|
||||
assert len(m.tables) == 1
|
||||
assert m.tables[0].section_key == "lieu_luong_va_cach_dung"
|
||||
assert "80 - 50" in m.tables[0].text
|
||||
assert "< 25 - 10" in m.tables[0].text
|
||||
assert "Không phối hợp với thuốc X." in m.sections["tuong_tac_thuoc"].text
|
||||
assert len({t.table_part_id for t in m.tables}) == len(m.tables)
|
||||
assert {t.continuation_group for t in m.tables} == {"p339_t0"}
|
||||
assert all(t.table_part_id.startswith("p339_t0@") for t in m.tables)
|
||||
assert all(t.quarantined for t in m.tables)
|
||||
|
||||
|
||||
def test_explicit_dose_adjustment_caption_reassigns_late_appendix_table():
|
||||
# CAPECITABIN p. 309: the PDF puts dose-adjustment tables after the trade
|
||||
# names and does not repeat the ordinary dosage section heading. Internal
|
||||
# block order can even emit a cell before the visually preceding caption.
|
||||
spans = [
|
||||
_span("CAPECITABIN", 309, 40.0, bold=True),
|
||||
_span("Tên chung quốc tế", 309, 50.0, bold=True),
|
||||
_span("Capecitabinum.", 309, 60.0),
|
||||
_span("Tên thương mại", 309, 70.0, bold=True),
|
||||
_span("Xeloda.", 309, 80.0),
|
||||
_span("Mức độ theo NCIC", 309, 120.0, block=5),
|
||||
_span("Bảng 3. Điều chỉnh liều do độc tính.", 309, 100.0),
|
||||
_span("Ngừng thuốc cho đến khi về mức 0.", 309, 140.0, block=5),
|
||||
]
|
||||
region = TableRegion("p309_t0", 309, (40.0, 115.0, 400.0, 180.0),
|
||||
3, 4, SHAPE_SIMPLE)
|
||||
m = list(assemble(spans, table_index=index_by_page([region])))[0]
|
||||
assert len(m.tables) == 1
|
||||
assert m.tables[0].section_key == "lieu_luong_va_cach_dung"
|
||||
dosage = m.sections["lieu_luong_va_cach_dung"].text
|
||||
assert dosage.count("Bảng 3. Điều chỉnh liều do độc tính.") == 1
|
||||
|
||||
@@ -53,6 +53,10 @@ def test_real_spelling_variants_found_in_the_book_all_match():
|
||||
"Hướng dẫn cách sử trí ADR": "huong_dan_xu_tri_adr",
|
||||
"Quá liều và xử lý": "qua_lieu_va_xu_tri",
|
||||
"Lọai thuốc": "loai_thuoc",
|
||||
"Tên chung quốc tế và mã ATC": "ten_chung_quoc_te",
|
||||
"Dạng bào chế và hàm lượng": "dang_thuoc_va_ham_luong",
|
||||
"Liều lượng và cách dùng giải độc tố uốn ván hấp phụ đơn giá":
|
||||
"lieu_luong_va_cach_dung",
|
||||
}
|
||||
for text, expected_key in cases.items():
|
||||
matched = match_section(text)
|
||||
|
||||
@@ -0,0 +1,72 @@
|
||||
from ingestion.validation.readiness import evaluate, evaluate_chunks
|
||||
|
||||
|
||||
def _count(monographs, gate_name):
|
||||
return next(g.count for g in evaluate(monographs) if g.name == gate_name)
|
||||
|
||||
|
||||
def test_readiness_rejects_a_part_without_source_span_provenance():
|
||||
corpus = [{
|
||||
"drug_id": "testdrug",
|
||||
"source_page_range": [100, 100],
|
||||
"sections": {
|
||||
"ma_atc": {
|
||||
"text": "N00AA00",
|
||||
"parts": [{"kind": "prose", "text": "N00AA00", "source_span_ids": []}],
|
||||
},
|
||||
},
|
||||
"tables": [],
|
||||
}]
|
||||
assert _count(corpus, "section_without_provenance") == 0
|
||||
assert _count(corpus, "part_without_source_span_ids") == 1
|
||||
|
||||
|
||||
def test_readiness_accepts_part_level_source_span_provenance():
|
||||
corpus = [{
|
||||
"drug_id": "testdrug",
|
||||
"source_page_range": [100, 100],
|
||||
"sections": {
|
||||
"ma_atc": {
|
||||
"text": "N00AA00",
|
||||
"parts": [{
|
||||
"kind": "prose", "text": "N00AA00",
|
||||
"source_span_ids": ["p100_b1_l0_s0"],
|
||||
}],
|
||||
},
|
||||
},
|
||||
"tables": [],
|
||||
}]
|
||||
assert _count(corpus, "part_without_source_span_ids") == 0
|
||||
|
||||
|
||||
def test_readiness_rejects_duplicate_physical_region_ids():
|
||||
corpus = [{
|
||||
"drug_id": "testdrug",
|
||||
"source_page_range": [100, 100],
|
||||
"sections": {},
|
||||
"tables": [
|
||||
{"table_id": "p100_t0", "quarantined": True},
|
||||
{"table_id": "p100_t0", "quarantined": True},
|
||||
],
|
||||
}]
|
||||
assert _count(corpus, "duplicate_table_id") == 1
|
||||
|
||||
|
||||
def test_chunk_readiness_requires_a_verified_printed_page_range():
|
||||
chunk = {
|
||||
"chunk_id": "drug__dose__0",
|
||||
"drug_id": "drug",
|
||||
"section_key": "dose",
|
||||
"chunk_kind": "prose",
|
||||
"text": "Dose.",
|
||||
"est_tokens": 2,
|
||||
"attachments": [],
|
||||
}
|
||||
gates = evaluate_chunks([], [chunk])
|
||||
missing = next(g for g in gates if g.name == "chunk_without_printed_page_range")
|
||||
assert missing.count == 1
|
||||
|
||||
chunk["printed_page_range"] = [101, 102]
|
||||
gates = evaluate_chunks([], [chunk])
|
||||
present = next(g for g in gates if g.name == "chunk_without_printed_page_range")
|
||||
assert present.count == 0
|
||||
Reference in New Issue
Block a user