72 lines
2.8 KiB
Python
72 lines
2.8 KiB
Python
"""Filesystem boundary for the chunk stage — kept out of the pure logic."""
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
from dataclasses import asdict
|
|
from pathlib import Path
|
|
from typing import Iterable, Iterator
|
|
|
|
from ..segment.models import Heading, Monograph, SectionSpan, TableBlock
|
|
from .models import SCHEMA_VERSION, Chunk
|
|
|
|
|
|
def read_monographs_jsonl(path: Path) -> Iterator[Monograph]:
|
|
with path.open(encoding="utf-8") as fh:
|
|
for line in fh:
|
|
line = line.strip()
|
|
if not line:
|
|
continue
|
|
raw = json.loads(line)
|
|
sections = {}
|
|
for key, s in raw.get("sections", {}).items():
|
|
h = s["heading"]
|
|
sections[key] = SectionSpan(
|
|
key=s["key"],
|
|
display_name=s["display_name"],
|
|
heading=Heading(
|
|
text=h["text"],
|
|
physical_page=h["physical_page"],
|
|
y0=h["y0"],
|
|
is_monograph_title=h["is_monograph_title"],
|
|
section_key=h.get("section_key"),
|
|
),
|
|
text=s["text"],
|
|
)
|
|
yield Monograph(
|
|
drug_id=raw["drug_id"],
|
|
drug_name=raw["drug_name"],
|
|
source_page_range=raw["source_page_range"],
|
|
sections=sections,
|
|
atc_codes=raw.get("atc_codes", []),
|
|
atc_stated_absent=raw.get("atc_stated_absent", False),
|
|
tables=[
|
|
TableBlock(
|
|
table_id=t["table_id"],
|
|
shape=t["shape"],
|
|
physical_page=t["physical_page"],
|
|
bbox=t["bbox"],
|
|
section_key=t.get("section_key"),
|
|
text=t.get("text", ""),
|
|
quarantined=t.get("quarantined", False),
|
|
table_part_id=t.get("table_part_id"),
|
|
continuation_group=t.get("continuation_group"),
|
|
source_span_ids=t.get("source_span_ids", []),
|
|
)
|
|
for t in raw.get("tables", [])
|
|
],
|
|
)
|
|
|
|
|
|
def write_chunks_jsonl(chunks: Iterable[Chunk], path: Path) -> int:
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
count = 0
|
|
with path.open("w", encoding="utf-8") as fh:
|
|
for chunk in chunks:
|
|
# ADR 0005 flagged the absence of a version and ADR 0006 made it
|
|
# necessary: the record now has two chunk kinds and an attachment
|
|
# list, so a consumer must be able to tell which shape it has.
|
|
record = {"schema_version": SCHEMA_VERSION, **asdict(chunk)}
|
|
fh.write(json.dumps(record, ensure_ascii=False) + "\n")
|
|
count += 1
|
|
return count
|