diff --git a/Makefile b/Makefile index aa1e86e9..39c714fb 100644 --- a/Makefile +++ b/Makefile @@ -57,6 +57,21 @@ GROBID_WAIT_INTERVAL ?= 5 BENCHMARK_PARSER_URL ?= $(SCIENCEBEAM_PARSER_URL) BENCHMARK_CONCURRENCY ?= 0 +TRAINING_DATA_OUTPUT ?= data/generated-training-data +TRAINING_DATA_NUM_WORKERS ?= 1 +# Per-document timeout in seconds; 0 disables. Skips outlier PDFs (e.g. 73-page, 38 MB) +# that cause the JATS aligner to run for many minutes. +TRAINING_DATA_DOCUMENT_TIMEOUT ?= 120 + +# Source training data (PDF + JATS XML) downloaded from the HF dataset. +# TRAINING_DATA_OUTPUT must point to a checkout of the output repo; create a +# symlink at data/generated-training-data or override the variable directly: +# make dev-generate-training-data TRAINING_DATA_OUTPUT=/path/to/output-repo +SOURCE_TRAINING_CONFIG ?= benchmarks/training-source.yml +SOURCE_TRAINING_DATA ?= data/source-training-data +SOURCE_TRAINING_MODE ?= smoke +SOURCE_TRAINING_SPLIT ?= train + SHOW_FIELD ?= SHOW_METHOD ?= edit_sim SHOW_CORPUS ?= biorxiv @@ -271,6 +286,32 @@ dev-benchmark-with-baselines: $(ARGS) +dev-fetch-training-source: + $(PYTHON) -m benchmarks.fetch_training_source_cli \ + --config $(SOURCE_TRAINING_CONFIG) \ + --mode $(SOURCE_TRAINING_MODE) \ + --split $(SOURCE_TRAINING_SPLIT) \ + --output-path $(SOURCE_TRAINING_DATA) + + +dev-generate-training-data: + @test -d "$(TRAINING_DATA_OUTPUT)" || { \ + echo "ERROR: TRAINING_DATA_OUTPUT='$(TRAINING_DATA_OUTPUT)' does not exist."; \ + echo " Clone the output repo and symlink it to data/generated-training-data,"; \ + echo " or pass TRAINING_DATA_OUTPUT=/path/to/repo on the command line."; \ + exit 1; } + TF_CPP_MIN_LOG_LEVEL=3 TF_ENABLE_ONEDNN_OPTS=0 \ + $(PYTHON) -m benchmarks.generate_training_data_cli \ + --config $(SOURCE_TRAINING_CONFIG) \ + --source-data $(SOURCE_TRAINING_DATA) \ + --output-path $(TRAINING_DATA_OUTPUT) \ + --split $(SOURCE_TRAINING_SPLIT) \ + --num-workers $(TRAINING_DATA_NUM_WORKERS) \ + --document-timeout $(TRAINING_DATA_DOCUMENT_TIMEOUT) \ + --debug \ + $(ARGS) + + docker-buildx-bake-build-all: docker buildx bake \ --file docker-bake.hcl \ diff --git a/benchmarks/fetch.py b/benchmarks/fetch.py index a8491f04..261fc534 100644 --- a/benchmarks/fetch.py +++ b/benchmarks/fetch.py @@ -59,7 +59,11 @@ def fetch_data( # pylint: disable=too-many-locals continue filename, id_column = _get_corpus_filename_and_id_column(corpus_cfg) - LOGGER.info("Fetching corpus %r (mode=%s, n=%d)", corpus, mode, sample_sizes[corpus]) + raw_n = sample_sizes[corpus] + LOGGER.info( + "Fetching corpus %r (mode=%s, n=%s)", corpus, mode, + raw_n if raw_n is not None else "all", + ) if local_root: parquet_path = str(Path(local_root) / filename) @@ -76,7 +80,7 @@ def fetch_data( # pylint: disable=too-many-locals # only the selected rows to avoid loading all PDFs into memory. pf = pq.ParquetFile(parquet_path) all_ids = pf.read(columns=[id_column]).column(id_column).to_pylist() - n = min(sample_sizes[corpus], len(all_ids)) + n = len(all_ids) if raw_n is None else min(raw_n, len(all_ids)) picked = _sample_indices(len(all_ids), n, seed) corpus_dir = data_dir / split / corpus @@ -135,8 +139,10 @@ def fetch_gold( # pylint: disable=too-many-locals continue filename, id_column = _get_corpus_filename_and_id_column(corpus_cfg) + raw_n = sample_sizes[corpus] LOGGER.info( - "Fetching gold for corpus %r (mode=%s, n=%d)", corpus, mode, sample_sizes[corpus] + "Fetching gold for corpus %r (mode=%s, n=%s)", corpus, mode, + raw_n if raw_n is not None else "all", ) if local_root: @@ -152,7 +158,7 @@ def fetch_gold( # pylint: disable=too-many-locals pf = pq.ParquetFile(parquet_path) all_ids = pf.read(columns=[id_column]).column(id_column).to_pylist() - n = min(sample_sizes[corpus], len(all_ids)) + n = len(all_ids) if raw_n is None else min(raw_n, len(all_ids)) picked = _sample_indices(len(all_ids), n, seed) corpus_dir = data_dir / split / corpus @@ -180,3 +186,23 @@ def fetch_gold( # pylint: disable=too-many-locals LOGGER.info("Corpus %r: materialised %d gold records to %s", corpus, n, corpus_dir) return records + + +def fetch_training_source( + cfg: Dict[str, Any], mode: str, split: str, data_dir: Path +) -> List[Dict[str, str]]: + """Fetch PDF + JATS XML for CC-BY corpora only. + + Reads ``cc_by_corpora`` from the config to determine which corpora are + permitted. Corpora absent from that list are silently skipped so that + the allow-list can be extended without changing call sites. + + Delegates to :func:`fetch_data` after building a filtered config. + """ + allowed = set(cfg.get("cc_by_corpora", [])) + filtered_sampling = { + m: {corpus: n for corpus, n in sizes.items() if corpus in allowed} + for m, sizes in cfg.get("sampling", {}).items() + } + filtered_cfg = {**cfg, "sampling": filtered_sampling} + return fetch_data(filtered_cfg, mode, split, data_dir) diff --git a/benchmarks/fetch_training_source_cli.py b/benchmarks/fetch_training_source_cli.py new file mode 100644 index 00000000..a7b868e1 --- /dev/null +++ b/benchmarks/fetch_training_source_cli.py @@ -0,0 +1,50 @@ +"""CLI: fetch CC-BY source data (PDF + JATS XML) for GROBID training data generation.""" +import argparse +import logging +from pathlib import Path + +import yaml + +from benchmarks.fetch import fetch_training_source + +LOGGER = logging.getLogger(__name__) + + +def main(argv=None): + parser = argparse.ArgumentParser( + description="Fetch CC-BY PDF + JATS XML pairs for GROBID training data generation." + ) + parser.add_argument( + "--config", + default="benchmarks/training-source.yml", + help="Path to training-source config YAML (default: benchmarks/training-source.yml)", + ) + parser.add_argument( + "--mode", + default="smoke", + help="Sampling mode defined in the config (e.g. smoke, small, full)", + ) + parser.add_argument( + "--split", + default="train", + help="Dataset split to fetch (default: train)", + ) + parser.add_argument( + "--output-path", + required=True, + help="Directory to write PDF and JATS XML files into", + ) + args = parser.parse_args(argv) + + logging.basicConfig( + level=logging.INFO, + format="%(asctime)s %(levelname)s %(name)s: %(message)s", + ) + + cfg = yaml.safe_load(Path(args.config).read_text(encoding="utf-8")) + records = fetch_training_source(cfg, args.mode, args.split, Path(args.output_path)) + LOGGER.info("Fetched %d records to %s", len(records), args.output_path) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/generate_training_data_cli.py b/benchmarks/generate_training_data_cli.py new file mode 100644 index 00000000..d1397444 --- /dev/null +++ b/benchmarks/generate_training_data_cli.py @@ -0,0 +1,91 @@ +"""CLI: generate GROBID training data for each CC-BY source corpus. + +Reads cc_by_corpora from training-source.yml and calls generate_data once per +corpus, writing output to ///. Any extra arguments +after -- are forwarded verbatim to generate_data. +""" +import argparse +import logging +import sys +from pathlib import Path + +import yaml + +from sciencebeam_parser.training.cli.generate_data import main as generate_data_main + +LOGGER = logging.getLogger(__name__) + + +def main(argv=None): + parser = argparse.ArgumentParser( + description=( + "Generate GROBID training data for all CC-BY corpora in the training-source config." + ), + # Allow forwarding unknown flags to generate_data + epilog="Any additional arguments are forwarded to generate_data.", + ) + parser.add_argument( + "--config", + default="benchmarks/training-source.yml", + help="Path to training-source config YAML", + ) + parser.add_argument( + "--source-data", + required=True, + help="Root directory of fetched source PDFs and JATS XML (e.g. data/source-training-data)", + ) + parser.add_argument( + "--output-path", + required=True, + help="Root directory of the output repo (e.g. data/generated-training-data)", + ) + parser.add_argument( + "--split", + default="train", + help="Dataset split subdirectory (default: train)", + ) + args, extra_argv = parser.parse_known_args(argv) + + logging.basicConfig( + level=logging.INFO, + format="%(asctime)s %(levelname)s %(name)s: %(message)s", + ) + + cfg = yaml.safe_load(Path(args.config).read_text(encoding="utf-8")) + corpora = cfg.get("cc_by_corpora", []) + if not corpora: + LOGGER.warning("No cc_by_corpora defined in %s; nothing to generate.", args.config) + sys.exit(0) + + errors = [] + for corpus in corpora: + corpus_source = Path(args.source_data) / args.split / corpus + if not corpus_source.exists(): + LOGGER.warning( + "Source directory not found for corpus %r, skipping: %s", corpus, corpus_source + ) + continue + + corpus_output = Path(args.output_path) / args.split / corpus + LOGGER.info("Generating training data for corpus %r -> %s", corpus, corpus_output) + + corpus_argv = [ + "--source-path", str(corpus_source / "*.pdf"), + "--source-xml-path", str(corpus_source / "*.jats.xml"), + "--output-path", str(corpus_output), + "--use-directory-structure", + *extra_argv, + ] + try: + generate_data_main(corpus_argv) + except Exception: # pylint: disable=broad-except + LOGGER.exception("Failed to generate training data for corpus %r", corpus) + errors.append(corpus) + + if errors: + LOGGER.error("Generation failed for corpora: %s", ", ".join(errors)) + sys.exit(1) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/tests/fetch_training_source_test.py b/benchmarks/tests/fetch_training_source_test.py new file mode 100644 index 00000000..6949b4a7 --- /dev/null +++ b/benchmarks/tests/fetch_training_source_test.py @@ -0,0 +1,118 @@ +from __future__ import annotations + +from pathlib import Path +from unittest.mock import MagicMock, patch + +from benchmarks.fetch import fetch_training_source + + +_BASE_CONFIG = { + "dataset": { + "repo_id": "org/repo", + "revision": "main", + "splits": { + "train": { + "ore": {"file": "ore/train.parquet", "id_column": "ppr_id"}, + "biorxiv": {"file": "biorxiv/train.parquet", "id_column": "ppr_id"}, + } + }, + }, + "cc_by_corpora": ["ore"], + "sampling": { + "smoke": {"ore": 2, "biorxiv": 2}, + "full": {"ore": None, "biorxiv": None}, + }, + "seeds": {"sample": 42}, +} + + +def _make_parquet_mock(ids: list) -> MagicMock: + pf = MagicMock() + id_col = MagicMock() + id_col.to_pylist.return_value = ids + pf.read.return_value.column.return_value = id_col + + batch = MagicMock() + batch.num_rows = len(ids) + + def _col(name): + col = MagicMock() + col.__getitem__ = lambda self, i: _cell(ids[i] if name == "ppr_id" else b"pdf") + return col + + def _cell(val): + m = MagicMock() + m.as_py.return_value = val + return m + + batch.column = _col + pf.iter_batches.return_value = [batch] + return pf + + +class TestFetchTrainingSource: + def test_only_fetches_cc_by_corpora(self, tmp_path: Path): + ore_pf = _make_parquet_mock(["ore1", "ore2", "ore3"]) + biorxiv_pf = _make_parquet_mock(["bx1", "bx2", "bx3"]) + + def _parquet_file(path): + return ore_pf if "ore" in path else biorxiv_pf + + with patch("benchmarks.fetch.hf_hub_download", return_value="fake.parquet"), \ + patch("benchmarks.fetch.pq.ParquetFile", side_effect=_parquet_file): + records = fetch_training_source(_BASE_CONFIG, "smoke", "train", tmp_path) + + corpora = {r["corpus"] for r in records} + assert "ore" in corpora + assert "biorxiv" not in corpora + + def test_full_mode_none_fetches_all_records(self, tmp_path: Path): + ids = [f"ore{i}" for i in range(5)] + pf = _make_parquet_mock(ids) + + with patch("benchmarks.fetch.hf_hub_download", return_value="fake.parquet"), \ + patch("benchmarks.fetch.pq.ParquetFile", return_value=pf): + records = fetch_training_source(_BASE_CONFIG, "full", "train", tmp_path) + + assert len(records) == len(ids) + + def test_empty_cc_by_corpora_fetches_nothing(self, tmp_path: Path): + cfg = {**_BASE_CONFIG, "cc_by_corpora": []} + pf = _make_parquet_mock(["id1", "id2"]) + with patch("benchmarks.fetch.hf_hub_download", return_value="fake.parquet"), \ + patch("benchmarks.fetch.pq.ParquetFile", return_value=pf): + records = fetch_training_source(cfg, "smoke", "train", tmp_path) + assert not records + + def test_writes_pdf_and_xml(self, tmp_path: Path): + pf = _make_parquet_mock(["ore1", "ore2"]) + # Provide xml column too + batch = MagicMock() + batch.num_rows = 2 + + def _col(name): + col = MagicMock() + if name == "ppr_id": + col.__getitem__ = lambda self, i: _cell(["ore1", "ore2"][i]) + elif name == "pdf": + col.__getitem__ = lambda self, i: _cell(b"pdfcontent") + else: + col.__getitem__ = lambda self, i: _cell("") + return col + + def _cell(val): + m = MagicMock() + m.as_py.return_value = val + return m + + batch.column = _col + pf.iter_batches.return_value = [batch] + + with patch("benchmarks.fetch.hf_hub_download", return_value="fake.parquet"), \ + patch("benchmarks.fetch.pq.ParquetFile", return_value=pf): + records = fetch_training_source(_BASE_CONFIG, "smoke", "train", tmp_path) + + assert len(records) == 2 + for r in records: + assert Path(r["pdf_path"]).exists() + assert Path(r["xml_path"]).exists() diff --git a/benchmarks/tests/generate_training_data_cli_test.py b/benchmarks/tests/generate_training_data_cli_test.py new file mode 100644 index 00000000..75afcede --- /dev/null +++ b/benchmarks/tests/generate_training_data_cli_test.py @@ -0,0 +1,200 @@ +from __future__ import annotations + +import textwrap +from pathlib import Path +from unittest.mock import patch + +import pytest + +from benchmarks.generate_training_data_cli import main + + +def _write_config(path: Path, corpora: list) -> Path: + config_file = path / "training-source.yml" + config_file.write_text( + textwrap.dedent(f"""\ + cc_by_corpora: {corpora!r} + """), + encoding="utf-8", + ) + return config_file + + +def _make_corpus_dir(source_root: Path, split: str, corpus: str) -> Path: + d = source_root / split / corpus + d.mkdir(parents=True) + return d + + +class TestGenerateTrainingDataCli: + def test_calls_generate_data_once_per_corpus(self, tmp_path: Path): + config = _write_config(tmp_path, ["ore", "scielo"]) + source = tmp_path / "source" + output = tmp_path / "output" + output.mkdir() + _make_corpus_dir(source, "train", "ore") + _make_corpus_dir(source, "train", "scielo") + + with patch("benchmarks.generate_training_data_cli.generate_data_main") as mock_gen: + main([ + "--config", str(config), + "--source-data", str(source), + "--output-path", str(output), + ]) + + assert mock_gen.call_count == 2 + called_corpora = [c.args[0][c.args[0].index("--output-path") + 1] + for c in mock_gen.call_args_list] + assert any("ore" in p for p in called_corpora) + assert any("scielo" in p for p in called_corpora) + + def test_source_and_output_paths_contain_split_and_corpus(self, tmp_path: Path): + config = _write_config(tmp_path, ["ore"]) + source = tmp_path / "source" + output = tmp_path / "output" + output.mkdir() + _make_corpus_dir(source, "train", "ore") + + with patch("benchmarks.generate_training_data_cli.generate_data_main") as mock_gen: + main([ + "--config", str(config), + "--source-data", str(source), + "--output-path", str(output), + "--split", "train", + ]) + + argv = mock_gen.call_args.args[0] + source_path = argv[argv.index("--source-path") + 1] + xml_path = argv[argv.index("--source-xml-path") + 1] + out_path = argv[argv.index("--output-path") + 1] + + assert source_path == str(source / "train" / "ore" / "*.pdf") + assert xml_path == str(source / "train" / "ore" / "*.jats.xml") + assert out_path == str(output / "train" / "ore") + + def test_use_directory_structure_always_forwarded(self, tmp_path: Path): + config = _write_config(tmp_path, ["ore"]) + source = tmp_path / "source" + output = tmp_path / "output" + output.mkdir() + _make_corpus_dir(source, "train", "ore") + + with patch("benchmarks.generate_training_data_cli.generate_data_main") as mock_gen: + main([ + "--config", str(config), + "--source-data", str(source), + "--output-path", str(output), + ]) + + argv = mock_gen.call_args.args[0] + assert "--use-directory-structure" in argv + + def test_extra_args_forwarded_to_generate_data(self, tmp_path: Path): + config = _write_config(tmp_path, ["ore"]) + source = tmp_path / "source" + output = tmp_path / "output" + output.mkdir() + _make_corpus_dir(source, "train", "ore") + + with patch("benchmarks.generate_training_data_cli.generate_data_main") as mock_gen: + main([ + "--config", str(config), + "--source-data", str(source), + "--output-path", str(output), + "--num-workers", "4", + "--document-timeout", "60", + "--debug", + ]) + + argv = mock_gen.call_args.args[0] + assert "--num-workers" in argv + assert "4" in argv + assert "--document-timeout" in argv + assert "60" in argv + assert "--debug" in argv + + def test_skips_corpus_with_missing_source_directory(self, tmp_path: Path): + config = _write_config(tmp_path, ["ore", "missing"]) + source = tmp_path / "source" + output = tmp_path / "output" + output.mkdir() + _make_corpus_dir(source, "train", "ore") + # "missing" corpus directory is intentionally not created + + with patch("benchmarks.generate_training_data_cli.generate_data_main") as mock_gen: + main([ + "--config", str(config), + "--source-data", str(source), + "--output-path", str(output), + ]) + + assert mock_gen.call_count == 1 + argv = mock_gen.call_args.args[0] + assert "ore" in argv[argv.index("--output-path") + 1] + + def test_empty_cc_by_corpora_exits_cleanly(self, tmp_path: Path): + config = _write_config(tmp_path, []) + source = tmp_path / "source" + output = tmp_path / "output" + output.mkdir() + + with patch("benchmarks.generate_training_data_cli.generate_data_main") as mock_gen: + with pytest.raises(SystemExit) as exc_info: + main([ + "--config", str(config), + "--source-data", str(source), + "--output-path", str(output), + ]) + + assert exc_info.value.code == 0 + mock_gen.assert_not_called() + + def test_corpus_failure_causes_exit_1(self, tmp_path: Path): + config = _write_config(tmp_path, ["ore"]) + source = tmp_path / "source" + output = tmp_path / "output" + output.mkdir() + _make_corpus_dir(source, "train", "ore") + + with patch( + "benchmarks.generate_training_data_cli.generate_data_main", + side_effect=RuntimeError("boom"), + ): + with pytest.raises(SystemExit) as exc_info: + main([ + "--config", str(config), + "--source-data", str(source), + "--output-path", str(output), + ]) + + assert exc_info.value.code == 1 + + def test_second_corpus_still_runs_after_first_fails(self, tmp_path: Path): + config = _write_config(tmp_path, ["ore", "scielo"]) + source = tmp_path / "source" + output = tmp_path / "output" + output.mkdir() + _make_corpus_dir(source, "train", "ore") + _make_corpus_dir(source, "train", "scielo") + + call_count = 0 + + def _side_effect(argv): + nonlocal call_count + call_count += 1 + if "ore" in argv[argv.index("--output-path") + 1]: + raise RuntimeError("ore failed") + + with patch( + "benchmarks.generate_training_data_cli.generate_data_main", + side_effect=_side_effect, + ): + with pytest.raises(SystemExit) as exc_info: + main([ + "--config", str(config), + "--source-data", str(source), + "--output-path", str(output), + ]) + + assert call_count == 2 + assert exc_info.value.code == 1 diff --git a/benchmarks/training-source.yml b/benchmarks/training-source.yml new file mode 100644 index 00000000..3b32bec2 --- /dev/null +++ b/benchmarks/training-source.yml @@ -0,0 +1,33 @@ +dataset: + repo_id: elifepathways/sciencebeam-v2-benchmarking + revision: main + splits: + train: + ore: + file: ore-jats/train-00000-of-00001.parquet + id_column: ppr_id + scielo_preprints-jats: + file: scielo-preprints-jats/train-00000-of-00001.parquet + id_column: ppr_id + +# Corpora confirmed CC-BY 4.0. Others are excluded until per-record licence +# data is available in the dataset. +cc_by_corpora: + - ore + - scielo_preprints-jats + +# Dataset row counts at revision main: +# ore: 199, scielo_preprints-jats: 1000 +sampling: + smoke: + ore: 5 + scielo_preprints-jats: 10 + small: + ore: 20 + scielo_preprints-jats: 50 + full: + ore: null # all records + scielo_preprints-jats: null + +seeds: + sample: 42 diff --git a/sciencebeam_parser/models/citation/training_data.py b/sciencebeam_parser/models/citation/training_data.py index 4219381c..649ea01d 100644 --- a/sciencebeam_parser/models/citation/training_data.py +++ b/sciencebeam_parser/models/citation/training_data.py @@ -38,6 +38,7 @@ '': ROOT_TRAINING_XML_ELEMENT_PATH + ['pubPlace'], '': ROOT_TRAINING_XML_ELEMENT_PATH + ['note[@type="report"]'], '': ROOT_TRAINING_XML_ELEMENT_PATH + ['ptr[@type="web"]'], + '': ROOT_TRAINING_XML_ELEMENT_PATH + ['idno[@type="DOI"]'], '': ROOT_TRAINING_XML_ELEMENT_PATH + ['idno'], '': ROOT_TRAINING_XML_ELEMENT_PATH + ['note'] } diff --git a/sciencebeam_parser/models/header/training_data.py b/sciencebeam_parser/models/header/training_data.py index 71b6c414..e8b7bb88 100644 --- a/sciencebeam_parser/models/header/training_data.py +++ b/sciencebeam_parser/models/header/training_data.py @@ -50,6 +50,13 @@ '': ROOT_TRAINING_XML_ELEMENT_PATH + ['byline', 'affiliation'] } +# Each new (B- prefix) resets to the front level so it gets its own +# wrapper rather than being appended to the author byline or a previous +# affiliation byline. +RESET_TRAINING_XML_ELEMENT_PATH_BY_LABEL = { + '': ROOT_TRAINING_XML_ELEMENT_PATH, +} + class HeaderTeiTrainingDataGenerator(AbstractTeiTrainingDataGenerator): DEFAULT_TEI_FILENAME_SUFFIX = '.header.tei.xml' @@ -60,6 +67,7 @@ def __init__(self): root_training_xml_element_path=ROOT_TRAINING_XML_ELEMENT_PATH, training_xml_element_path_by_label=TRAINING_XML_ELEMENT_PATH_BY_LABEL, element_maker=TEI_E, + reset_training_xml_element_path_by_label=RESET_TRAINING_XML_ELEMENT_PATH_BY_LABEL, default_tei_filename_suffix=( HeaderTeiTrainingDataGenerator.DEFAULT_TEI_FILENAME_SUFFIX ), diff --git a/sciencebeam_parser/training/cli/generate_data.py b/sciencebeam_parser/training/cli/generate_data.py index 4b2fb27f..01cbfebb 100644 --- a/sciencebeam_parser/training/cli/generate_data.py +++ b/sciencebeam_parser/training/cli/generate_data.py @@ -1,9 +1,13 @@ +# pylint: disable=too-many-lines from abc import ABC, abstractmethod import argparse import logging import os +import multiprocessing +import time + from dataclasses import dataclass, field -from typing import Dict, Iterable, List, NamedTuple, Optional, Sequence +from typing import Callable, Dict, Iterable, List, NamedTuple, Optional, Sequence, Tuple from lxml import etree @@ -42,6 +46,17 @@ from sciencebeam_parser.config.config import AppConfig from sciencebeam_parser.app.parser import ScienceBeamParser from sciencebeam_parser.utils.media_types import MediaTypes +from sciencebeam_parser.training.jats.annotated_document import JatsAnnotatedLayoutDocument +from sciencebeam_parser.training.jats.field_vocab import ( + AFF_LABEL_BY_SUB_FIELD, + CITATION_LABEL_BY_SUB_FIELD, + FULLTEXT_LABEL_BY_FIELD, + HEADER_LABEL_BY_FIELD, + JatsSubFieldNames, +) +from sciencebeam_parser.training.jats.field_extractor import JatsFieldExtractor +from sciencebeam_parser.training.jats.aligner import LayoutDocumentJatsAligner +from sciencebeam_parser.training.jats.segmentation import SegmentationLabelDeriver LOGGER = logging.getLogger(__name__) @@ -93,6 +108,43 @@ def parse_args(argv: Optional[List[str]] = None) -> argparse.Namespace: action='store_true', help='Enable debug logging' ) + parser.add_argument( + '--source-xml-path', + type=str, + required=False, + help='Glob pattern to JATS XML files; matched to PDFs by filename stem' + ) + parser.add_argument( + '--required-fields', + type=str, + nargs='*', + default=[], + help='JATS field names that must be present AND aligned; skip document if any are missing' + ) + parser.add_argument( + '--require-matching-fields', + type=str, + nargs='*', + default=[], + help='JATS field names that must align if present in JATS XML' + ) + parser.add_argument( + '--num-workers', + type=int, + default=1, + help='Number of parallel worker processes (default: 1)' + ) + parser.add_argument( + '--document-timeout', + type=int, + default=0, + metavar='SECONDS', + help=( + 'Per-document time limit in seconds (0 = no limit). ' + 'Documents that exceed this limit are skipped with a warning. ' + 'Single-worker mode uses SIGALRM; multi-worker mode uses future timeout.' + ) + ) return parser.parse_args(argv) @@ -175,6 +227,8 @@ class TrainingDataDocumentContext(NamedTuple): use_directory_structure: bool model_result_cache: ModelResultCache gzip_enabled: bool + jats_annotated_document: Optional[JatsAnnotatedLayoutDocument] = None + jats_segmentation_labels: Optional[Dict[int, str]] = None @property def source_name(self) -> str: @@ -286,10 +340,36 @@ def get_labeled_layout_tokens_for_model_and_layout_document( return labeled_layout_tokens_list[0] +def _get_jats_segmentation_label_result( + layout_document: LayoutDocument, + jats_segmentation_labels: Dict[int, str], +) -> LayoutDocumentLabelResult: + """Build a LayoutDocumentLabelResult from JATS-derived per-line segmentation labels.""" + layout_model_labels = [ + LayoutModelLabel( + label=jats_segmentation_labels.get(id(line), ''), + label_token_text=line.text, + layout_line=line, + layout_token=None, + ) + for block in layout_document.iter_all_blocks() + for line in block.lines + ] + return LayoutDocumentLabelResult( + layout_document=layout_document, + layout_model_label_iterable=layout_model_labels, + ) + + def get_segmentation_label_result( layout_document: LayoutDocument, document_context: TrainingDataDocumentContext ) -> LayoutDocumentLabelResult: + if document_context.jats_segmentation_labels is not None: + return _get_jats_segmentation_label_result( + layout_document=layout_document, + jats_segmentation_labels=document_context.jats_segmentation_labels, + ) segmentation_label_model_data_lists = list( iter_labeled_model_data_list_for_model_and_layout_documents( model=document_context.fulltext_models.segmentation_model, @@ -305,6 +385,24 @@ def get_segmentation_label_result( ) +JatsLabelFn = Callable[ + [JatsAnnotatedLayoutDocument, Dict[int, str], LayoutModelData], + Optional[str] +] + + +def _apply_jats_labels_to_model_data_list( + model_data_list: Sequence[LayoutModelData], + annotated: JatsAnnotatedLayoutDocument, + jats_seg_labels: Dict[int, str], + label_fn: JatsLabelFn, +) -> Sequence[LabeledLayoutModelData]: + return [ + LabeledLayoutModelData.from_model_data(md, label=label_fn(annotated, jats_seg_labels, md)) + for md in model_data_list + ] + + class AbstractModelTrainingDataGenerator(ABC): def get_pre_file_path_suffix(self) -> str: return '' @@ -413,6 +511,10 @@ def iter_model_layout_documents( ) -> Iterable[LayoutDocument]: pass + def get_jats_label_fn(self) -> Optional[JatsLabelFn]: + """Return a JATS label function, or None to skip JATS labeling for this model.""" + return None + def iter_model_data_list( self, layout_document: LayoutDocument, @@ -423,6 +525,23 @@ def iter_model_data_list( layout_document, document_context=document_context )) + annotated = document_context.jats_annotated_document + jats_seg_labels = document_context.jats_segmentation_labels + jats_label_fn = self.get_jats_label_fn() # pylint: disable=assignment-from-none + if annotated is not None and jats_seg_labels is not None and jats_label_fn is not None: + unlabeled_lists = list( + iter_unlabeled_model_data_list_for_model_and_layout_documents( + model=model, + model_layout_documents=model_layout_documents, + document_context=document_context, + ) + ) + return [ + _apply_jats_labels_to_model_data_list( + mdl, annotated, jats_seg_labels, jats_label_fn + ) + for mdl in unlabeled_lists + ] return iter_model_data_list_for_model_and_layout_documents( model=model, model_layout_documents=model_layout_documents, @@ -441,11 +560,64 @@ def iter_model_layout_documents( ) -> Iterable[LayoutDocument]: return [layout_document] + def get_jats_label_fn(self) -> Optional[JatsLabelFn]: + def fn( + _annotated: JatsAnnotatedLayoutDocument, + seg_labels: Dict[int, str], + md: LayoutModelData, + ) -> Optional[str]: + return seg_labels.get(id(md.layout_line)) if md.layout_line else None + return fn + class HeaderModelTrainingDataGenerator(AbstractDocumentModelTrainingDataGenerator): def get_main_model(self, document_context: TrainingDataDocumentContext) -> Model: return document_context.fulltext_models.header_model + def get_jats_label_fn(self) -> Optional[JatsLabelFn]: + # Stateful closure: emit B-/I- IOB prefix so the TEI generator can create + # separate blocks for each JATS element. + # Address sub-fields (city, region, postcode, country, bulk addr range) are + # mapped to
instead of . + _HEADER_ADDRESS_SUB_FIELDS = frozenset({ + JatsSubFieldNames.AUTHOR_AFF_ADDR, + JatsSubFieldNames.AUTHOR_AFF_CITY, + JatsSubFieldNames.AUTHOR_AFF_POSTCODE, + JatsSubFieldNames.AUTHOR_AFF_REGION, + JatsSubFieldNames.AUTHOR_AFF_COUNTRY, + }) + prev_label_instance: Optional[Tuple[str, int]] = None + + def fn( + annotated: JatsAnnotatedLayoutDocument, + _seg_labels: Dict[int, str], + md: LayoutModelData, + ) -> Optional[str]: + nonlocal prev_label_instance + token = md.layout_token + if not token: + prev_label_instance = None + return None + field_name = annotated.get_token_field(token) + if not field_name: + prev_label_instance = None + return None + sub_field_name = annotated.get_token_sub_field(token) + if sub_field_name in _HEADER_ADDRESS_SUB_FIELDS: + label: Optional[str] = '
' + else: + label = HEADER_LABEL_BY_FIELD.get(field_name) + if label is None: + prev_label_instance = None + return None + instance_id = annotated.get_token_instance(token) + label_instance = (label, instance_id) + prefix = 'B' if label_instance != prev_label_instance else 'I' + prev_label_instance = label_instance + return f'{prefix}-{label}' + + return fn + def iter_model_layout_documents( self, layout_document: LayoutDocument, @@ -469,6 +641,19 @@ class AffiliationAddressModelTrainingDataGenerator(AbstractDocumentModelTraining def get_main_model(self, document_context: TrainingDataDocumentContext) -> Model: return document_context.fulltext_models.affiliation_address_model + def get_jats_label_fn(self) -> Optional[JatsLabelFn]: + def fn( + annotated: JatsAnnotatedLayoutDocument, + _seg_labels: Dict[int, str], + md: LayoutModelData, + ) -> Optional[str]: + token = md.layout_token + if not token: + return None + sub_field = annotated.get_token_sub_field(token) + return AFF_LABEL_BY_SUB_FIELD.get(sub_field or '') if sub_field else None + return fn + def iter_model_layout_documents( self, layout_document: LayoutDocument, @@ -648,6 +833,19 @@ class FullTextModelTrainingDataGenerator(AbstractDocumentModelTrainingDataGenera def get_main_model(self, document_context: TrainingDataDocumentContext) -> Model: return document_context.fulltext_models.fulltext_model + def get_jats_label_fn(self) -> Optional[JatsLabelFn]: + def fn( + annotated: JatsAnnotatedLayoutDocument, + _seg_labels: Dict[int, str], + md: LayoutModelData, + ) -> Optional[str]: + token = md.layout_token + if not token: + return None + field_name = annotated.get_token_field(token) + return FULLTEXT_LABEL_BY_FIELD.get(field_name or '') if field_name else None + return fn + def iter_model_layout_documents( self, layout_document: LayoutDocument, @@ -772,6 +970,19 @@ class CitationModelTrainingDataGenerator(AbstractDocumentModelTrainingDataGenera def get_main_model(self, document_context: TrainingDataDocumentContext) -> Model: return document_context.fulltext_models.citation_model + def get_jats_label_fn(self) -> Optional[JatsLabelFn]: + def fn( + annotated: JatsAnnotatedLayoutDocument, + _seg_labels: Dict[int, str], + md: LayoutModelData, + ) -> Optional[str]: + token = md.layout_token + if not token: + return None + sub_field = annotated.get_token_sub_field(token) + return CITATION_LABEL_BY_SUB_FIELD.get(sub_field or '') if sub_field else None + return fn + def iter_model_layout_documents( self, layout_document: LayoutDocument, @@ -812,6 +1023,21 @@ def iter_model_layout_documents( ] +def _build_jats_annotations( + layout_document: LayoutDocument, + jats_xml_filename: str, +) -> Optional[JatsAnnotatedLayoutDocument]: + try: + with auto_download_input_file(jats_xml_filename, auto_decompress=True) as local_xml: + root = etree.parse(local_xml).getroot() + except Exception: # pylint: disable=broad-except + LOGGER.warning('Failed to load JATS XML: %r', jats_xml_filename, exc_info=True) + return None + field_values = list(JatsFieldExtractor().iter_field_values(root)) + LOGGER.debug('JATS field values count: %d', len(field_values)) + return LayoutDocumentJatsAligner().align(layout_document, field_values) + + def generate_training_data_for_layout_document( layout_document: LayoutDocument, *, @@ -821,9 +1047,21 @@ def generate_training_data_for_layout_document( fulltext_models: FullTextModels, use_model: bool, use_directory_structure: bool, - gzip_enabled: bool = False + gzip_enabled: bool = False, + jats_xml_filename: Optional[str] = None, ): model_result_cache = ModelResultCache() + jats_annotated: Optional[JatsAnnotatedLayoutDocument] = None + jats_seg_labels: Optional[Dict[int, str]] = None + if jats_xml_filename: + jats_annotated = _build_jats_annotations(layout_document, jats_xml_filename) + if jats_annotated: + jats_seg_labels = SegmentationLabelDeriver().derive_labels( + layout_document, jats_annotated + ) + LOGGER.debug( + 'JATS coverage ratio: %.2f', jats_annotated.coverage_ratio() + ) document_context = TrainingDataDocumentContext( output_path=output_path, source_filename=source_filename, @@ -832,7 +1070,9 @@ def generate_training_data_for_layout_document( use_model=use_model, use_directory_structure=use_directory_structure, model_result_cache=model_result_cache, - gzip_enabled=gzip_enabled + gzip_enabled=gzip_enabled, + jats_annotated_document=jats_annotated, + jats_segmentation_labels=jats_seg_labels, ) training_data_generators = [ SegmentationModelTrainingDataGenerator(), @@ -867,6 +1107,23 @@ def get_layout_document_for_source_filename( return layout_document +def _find_jats_xml_for_source( + source_filename: str, + xml_file_list: Sequence[str], +) -> Optional[str]: + """Find the XML file whose stem matches the PDF stem, stripping compound extensions.""" + source_stem = os.path.splitext(os.path.basename(source_filename))[0] + for xml_filename in xml_file_list: + xml_stem = os.path.basename(xml_filename) + while True: + xml_stem, ext = os.path.splitext(xml_stem) + if xml_stem == source_stem: + return xml_filename + if not ext: + break + return None + + def generate_training_data_for_source_filename( source_filename: str, *, @@ -874,13 +1131,21 @@ def generate_training_data_for_source_filename( sciencebeam_parser: ScienceBeamParser, use_model: bool, use_directory_structure: bool, - gzip_enabled: bool + gzip_enabled: bool, + xml_file_list: Optional[Sequence[str]] = None, ): LOGGER.debug('use_model: %r', use_model) layout_document = get_layout_document_for_source_filename( source_filename, sciencebeam_parser=sciencebeam_parser ) + jats_xml_filename: Optional[str] = None + if xml_file_list: + jats_xml_filename = _find_jats_xml_for_source(source_filename, xml_file_list) + if jats_xml_filename: + LOGGER.info('Using JATS XML: %r', jats_xml_filename) + else: + LOGGER.warning('No matching JATS XML found for: %r', source_filename) generate_training_data_for_layout_document( layout_document=layout_document, output_path=output_path, @@ -891,7 +1156,8 @@ def generate_training_data_for_source_filename( fulltext_models=sciencebeam_parser.fulltext_models, use_model=use_model, use_directory_structure=use_directory_structure, - gzip_enabled=gzip_enabled + gzip_enabled=gzip_enabled, + jats_xml_filename=jats_xml_filename, ) @@ -904,6 +1170,184 @@ def get_source_file_list_or_fail( return source_file_list +def _format_eta(seconds: float) -> str: + if seconds >= 3600: + return f'{seconds / 3600:.1f}h' + if seconds >= 60: + return f'{int(seconds) // 60}m{int(seconds) % 60:02d}s' + return f'{seconds:.0f}s' + + +class _Progress: + def __init__(self, total: int) -> None: + self.total = total + self.n_ok = 0 + self.n_err = 0 + self._t_start = time.monotonic() + + @property + def completed(self) -> int: + return self.n_ok + self.n_err + + def record(self, source_filename: str, ok: bool, elapsed_s: float) -> None: + if ok: + self.n_ok += 1 + else: + self.n_err += 1 + elapsed_total = time.monotonic() - self._t_start + done = self.completed + rate = done / elapsed_total if elapsed_total > 0 else 0.0 + remaining = self.total - done + eta = _format_eta(remaining / rate) if rate > 0 else '?' + status = 'ok' if ok else 'err' + LOGGER.info( + '[%d/%d] %s %s %.1fs | %.2f doc/s | ~%s left', + done, self.total, + os.path.basename(source_filename), status, elapsed_s, rate, eta, + ) + + +# Module-level worker state, initialised once per worker process. +_worker_sciencebeam_parser: Optional[ScienceBeamParser] = None + + +def _worker_init() -> None: + global _worker_sciencebeam_parser # pylint: disable=global-statement + config = AppConfig.load_yaml(DEFAULT_CONFIG_FILE) + _worker_sciencebeam_parser = ScienceBeamParser.from_config(config) + + +def _worker_process(kwargs: dict) -> bool: + assert _worker_sciencebeam_parser is not None + try: + generate_training_data_for_source_filename( + kwargs['source_filename'], + output_path=kwargs['output_path'], + sciencebeam_parser=_worker_sciencebeam_parser, + use_model=kwargs['use_model'], + use_directory_structure=kwargs['use_directory_structure'], + gzip_enabled=kwargs['gzip_enabled'], + xml_file_list=kwargs['xml_file_list'], + ) + return True + except Exception: # pylint: disable=broad-except + LOGGER.exception('Failed to process %r', kwargs['source_filename']) + return False + + +def _run_serial( + source_file_list: Sequence[str], + output_path: str, + args: argparse.Namespace, + xml_file_list: Optional[Sequence[str]], + progress: '_Progress', + document_timeout: int = 0, +) -> None: + """Run documents sequentially. + + When document_timeout == 0: runs inline in the current process. This + avoids subprocess-spawn overhead and works correctly with test mocks on + platforms that use the 'spawn' start method (e.g. macOS). + + When document_timeout > 0: uses multiprocessing.Pool so that + pool.terminate() can kill a worker stuck inside a C extension + (signal.alarm cannot interrupt C code). The pool is recreated after + each timeout so subsequent documents can run. + """ + common_kwargs = { + 'output_path': output_path, + 'use_model': args.use_model, + 'use_directory_structure': args.use_directory_structure, + 'gzip_enabled': args.gzip, + 'xml_file_list': xml_file_list, + } + + if document_timeout == 0: + # No timeout needed — run inline without spawning a subprocess. + _worker_init() + for source_filename in source_file_list: + kwargs = {'source_filename': source_filename, **common_kwargs} + t0 = time.monotonic() + ok = _worker_process(kwargs) + progress.record(source_filename, ok=ok, elapsed_s=time.monotonic() - t0) + return + + pool = multiprocessing.Pool(1, initializer=_worker_init) # pylint: disable=consider-using-with + try: + for source_filename in source_file_list: + kwargs = {'source_filename': source_filename, **common_kwargs} + t0 = time.monotonic() + async_result = pool.apply_async(_worker_process, (kwargs,)) + try: + ok = async_result.get(timeout=document_timeout) + except multiprocessing.TimeoutError: + LOGGER.warning( + 'Document exceeded %ds timeout, skipping: %r', + document_timeout, source_filename, + ) + pool.terminate() + pool.join() + # pylint: disable-next=consider-using-with + pool = multiprocessing.Pool(1, initializer=_worker_init) + ok = False + except Exception: # pylint: disable=broad-except + LOGGER.exception('Failed to process %r', source_filename) + ok = False + progress.record(source_filename, ok=ok, elapsed_s=time.monotonic() - t0) + finally: + pool.close() + pool.join() + + +def _run_parallel_workers( + source_file_list: Sequence[str], + output_path: str, + args: argparse.Namespace, + xml_file_list: Optional[Sequence[str]], + progress: '_Progress', + num_workers: int, + document_timeout: int = 0, +) -> None: + """Process documents in parallel using a multiprocessing.Pool. + + All documents are submitted upfront so workers stay busy. Results are + collected in submission order; pool.terminate() at the end kills any worker + that is still stuck in a C extension after its timeout. + """ + common_kwargs = { + 'output_path': output_path, + 'use_model': args.use_model, + 'use_directory_structure': args.use_directory_structure, + 'gzip_enabled': args.gzip, + 'xml_file_list': xml_file_list, + } + # pylint: disable-next=consider-using-with + pool = multiprocessing.Pool(num_workers, initializer=_worker_init) + timeout_arg = document_timeout if document_timeout > 0 else None + work = [ + (sf, pool.apply_async(_worker_process, ({'source_filename': sf, **common_kwargs},))) + for sf in source_file_list + ] + try: + for source_filename, async_result in work: + t0 = time.monotonic() + try: + ok = async_result.get(timeout=timeout_arg) + except multiprocessing.TimeoutError: + LOGGER.warning( + 'Document exceeded %ds timeout, skipping: %r', + document_timeout, source_filename, + ) + ok = False + except Exception: # pylint: disable=broad-except + LOGGER.exception('Failed to process %r', source_filename) + ok = False + progress.record(source_filename, ok=ok, elapsed_s=time.monotonic() - t0) + finally: + pool.terminate() + pool.join() + + def run(args: argparse.Namespace): LOGGER.info('args: %r', args) source_file_list = get_source_file_list_or_fail(args.source_path) @@ -911,30 +1355,43 @@ def run(args: argparse.Namespace): source_file_list = source_file_list[:args.limit] LOGGER.info('source files: %d', len(source_file_list)) output_path = args.output_path - config = AppConfig.load_yaml( - DEFAULT_CONFIG_FILE - ) - sciencebeam_parser = ScienceBeamParser.from_config(config) LOGGER.info('output_path: %r', output_path) + xml_file_list: Optional[Sequence[str]] = None + if args.source_xml_path: + xml_file_list = list(glob(args.source_xml_path)) + LOGGER.info('JATS XML files: %d', len(xml_file_list)) # Note: creating the directory may not be necessary, but provides early feedback makedirs(output_path, exist_ok=True) - for source_filename in source_file_list: - generate_training_data_for_source_filename( - source_filename, - output_path=output_path, - sciencebeam_parser=sciencebeam_parser, - use_model=args.use_model, - use_directory_structure=args.use_directory_structure, - gzip_enabled=args.gzip + total = len(source_file_list) + progress = _Progress(total) + num_workers = getattr(args, 'num_workers', 1) + document_timeout: int = getattr(args, 'document_timeout', 0) + + if num_workers > 1: + _run_parallel_workers( + source_file_list, output_path, args, xml_file_list, progress, num_workers, + document_timeout=document_timeout, ) + else: + _run_serial( + source_file_list, output_path, args, xml_file_list, progress, + document_timeout=document_timeout, + ) + + if progress.n_err: + LOGGER.warning('%d/%d documents failed', progress.n_err, total) + LOGGER.info('Done. Processed %d/%d documents.', progress.n_ok, total) def main(argv: Optional[List[str]] = None): LOGGER.debug('argv: %r', argv) args = parse_args(argv) if args.debug: - for name in [__name__, 'sciencebeam_parser', 'sciencebeam_trainer_delft']: - logging.getLogger(name).setLevel('DEBUG') + # Only enable DEBUG for the training CLI itself. Library loggers + # (model inference, aligner per-field traces) stay at INFO to avoid + # flooding the output with thousands of internal messages. + logging.getLogger(__name__).setLevel('DEBUG') + logging.getLogger('sciencebeam_parser.training').setLevel('DEBUG') run(args) diff --git a/sciencebeam_parser/training/jats/__init__.py b/sciencebeam_parser/training/jats/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/sciencebeam_parser/training/jats/aligner.py b/sciencebeam_parser/training/jats/aligner.py new file mode 100644 index 00000000..85bb5bb7 --- /dev/null +++ b/sciencebeam_parser/training/jats/aligner.py @@ -0,0 +1,538 @@ +import logging +from dataclasses import dataclass +from typing import Dict, FrozenSet, List, Optional, Set, Tuple + +from sciencebeam_alignment.align import LocalSequenceMatcher, SimpleScoring + +from sciencebeam_parser.document.layout_document import LayoutDocument, LayoutToken +from sciencebeam_parser.training.jats.annotated_document import JatsAnnotatedLayoutDocument +from sciencebeam_parser.training.jats.field_extractor import JatsFieldValue +from sciencebeam_parser.training.jats.field_vocab import JatsFieldNames +from sciencebeam_parser.training.jats.text_normalizer import normalize_for_alignment + + +LOGGER = logging.getLogger(__name__) + +_NO_TOKEN_INDEX = -1 + +# Fields whose matches establish the document region boundary (end of front matter). +# Body content is searched from the end of the last anchor match, so that front-matter +# fields which false-match citations late in the document do not push the search start +# past the actual body position. +_ANCHOR_FIELDS: FrozenSet[str] = frozenset({ + JatsFieldNames.TITLE, + JatsFieldNames.ABSTRACT, +}) + +# The abstract match establishes the end of the front-matter region. All non-anchor, +# non-body fields (authors, affiliations, keywords) are then confined to +# [0, abstract_end + _FRONT_MATTER_BUFFER] so they cannot false-match citations that +# appear later in the document. Authors physically precede the abstract in most PDFs +# but follow it in JATS ordering, so without this constraint they would be searched +# from last_match_end (≈abstract end) and miss their true page-1 position. +_FRONT_MATTER_END_FIELDS: FrozenSet[str] = frozenset({JatsFieldNames.ABSTRACT}) +_FRONT_MATTER_BUFFER = 2000 + +# When the "Keywords" section header is matched, individual keyword values are searched +# from that position forward rather than from position 0. Without this, short common +# keywords ("confidence", "Bayesian") false-match in the title or abstract. +_KEYWORDS_SECTION_ANCHOR_FIELDS: FrozenSet[str] = frozenset({JatsFieldNames.KEYWORDS_TITLE}) +_KEYWORDS_FIELDS: FrozenSet[str] = frozenset({JatsFieldNames.KEYWORDS}) + +# Fields that appear after the front matter region. They search from the body floor +# (end of last anchor match) rather than from the global last_match_end. +_BODY_CONTENT_FIELDS: FrozenSet[str] = frozenset({ + JatsFieldNames.BODY_SECTION_TITLE, + JatsFieldNames.BODY_SECTION_PARAGRAPH, + JatsFieldNames.BODY_FIGURE, + JatsFieldNames.BODY_TABLE, + JatsFieldNames.ACK_SECTION_TITLE, + JatsFieldNames.ACK_SECTION_PARAGRAPH, + JatsFieldNames.APPENDIX_GROUP_TITLE, + JatsFieldNames.APPENDIX, + JatsFieldNames.BACK_SECTION_TITLE, + JatsFieldNames.BACK_SECTION_PARAGRAPH, +}) + +# Reference fields use a dedicated floor so that appendix/body content matched +# after the reference section cannot push body_content_end past the references. +# reference_list_title anchors the search; references then advance reference_floor +# incrementally. Without this separation, body content from mathematical appendices +# (which may physically appear after references in the PDF) would advance +# body_content_end past all reference positions. +_REFERENCE_ANCHOR_FIELDS: FrozenSet[str] = frozenset({ + JatsFieldNames.REFERENCE_LIST_TITLE, +}) +_REFERENCE_FIELDS: FrozenSet[str] = frozenset({ + JatsFieldNames.REFERENCE, +}) + +# Fields that appear entirely after the main body and reference sections +# (e.g. ORE peer-review sub-articles). Searched from last_match_end rather than +# from the front-matter window, preventing author-response text (which often quotes +# the paper verbatim) from overwriting body / figure tokens via the global fallback. +_POST_BODY_FIELDS: FrozenSet[str] = frozenset({ + JatsFieldNames.SUB_ARTICLE, +}) + +# Smith-Waterman scoring: match=2, mismatch=-1, gap=-1 +_SCORING = SimpleScoring(match_score=2, mismatch_score=-1, gap_score=-1) + +# Window size limits to keep LocalSequenceMatcher fast (O(window * needle)) +_DEFAULT_MIN_WINDOW = 2000 +_WINDOW_NEEDLE_MULTIPLIER = 6 + +# Sub-field containment buffer: search sub-fields this many chars beyond the parent's +# matched range. Keeps short sub-field values (e.g. "USA", "2020") from matching +# identical text elsewhere in the document. +_SUB_FIELD_PARENT_BUFFER = 200 +_SUB_FIELD_PARENT_PRE_BUFFER = 0 + +# Anchor+chain labelling strategy: +# Smith-Waterman produces many tiny (1–4 char) matching blocks while traversing +# interleaved sidebar content. Those blocks must not cause sidebar tokens to be +# labelled. We use two constants: +# +# _MIN_ANCHOR_BLOCK_SIZE: a block is an "anchor" only if it is at least this +# many characters long. Sidebar words are almost never exact multi-word +# substrings of the abstract/keywords needle, so their SW blocks stay small. +# +# _MAX_HAYSTACK_GAP_TO_FILL: a small (non-anchor) block is included only if it +# starts within this many characters of the previous *included* block end. +# This fills legitimate intra-field gaps (e.g. the 2-char comma gap between +# "confidence, bayesian," and "ddm") without re-entering a sidebar whose +# last anchor lies hundreds of chars earlier. +# +# Fallback: when a field value produces NO anchor blocks at all (the entire text +# is shorter than _MIN_ANCHOR_BLOCK_SIZE), every block is labelled so short +# fields are never silently dropped. +_MIN_ANCHOR_BLOCK_SIZE = 5 +_MAX_HAYSTACK_GAP_TO_FILL = 3 + +# Type alias for the return value of _fuzzy_match_field_value: +# (abs_start, abs_end, [(block_start, block_end), ...]) +_MatchResult = Tuple[int, int, List[Tuple[int, int]]] + + +@dataclass +class AlignmentConfig: + threshold: float = 0.8 + max_window: int = 8000 + + +class _TokenIndex: + """Flat character-level haystack built from all layout tokens. + + Tracks which character offset belongs to which token so that a match + range [a, b) in the haystack can be mapped back to a set of tokens. + + Line-break hyphens are removed and the two word halves concatenated so + that tokens ["hyphen", "-"] at end of line followed by ["ation"] at the + start of the next line appear as "hyphenation" in the haystack, matching + the unhyphenated form found in the JATS source text. + """ + + def __init__( + self, + tokens: List[LayoutToken], + skip_tokens: Optional[Set[int]] = None, + no_space_after: Optional[Set[int]] = None, + ) -> None: + self.tokens = tokens + if skip_tokens is None: + skip_tokens = set() + if no_space_after is None: + no_space_after = set() + parts: List[str] = [] + token_index_at: List[int] = [] + + for tok_idx, token in enumerate(tokens): + if tok_idx in skip_tokens: + continue + norm = normalize_for_alignment(token.text) + if not norm: + continue + for _ in norm: + token_index_at.append(tok_idx) + parts.append(norm) + if tok_idx not in no_space_after: + parts.append(' ') + token_index_at.append(_NO_TOKEN_INDEX) + + self.haystack = ''.join(parts) + self._token_index_at = token_index_at + self._skip_tokens = skip_tokens + + def tokens_in_range(self, start: int, end: int) -> List[LayoutToken]: + seen: Set[int] = set() + result_indices: List[int] = [] + for i in range(start, min(end, len(self._token_index_at))): + tok_idx = self._token_index_at[i] + if tok_idx != _NO_TOKEN_INDEX and tok_idx not in seen: + seen.add(tok_idx) + result_indices.append(tok_idx) + # Include bare end-of-line hyphen tokens (skip_tokens) whose preceding + # word token was collected. These hyphens are invisible in the haystack + # but are physically part of the hyphenated word and should carry the + # same label as the surrounding tokens. + if self._skip_tokens: + filled: List[int] = [] + added_skips: Set[int] = set() + for tok_idx in result_indices: + filled.append(tok_idx) + next_idx = tok_idx + 1 + if next_idx in self._skip_tokens and next_idx not in added_skips: + filled.append(next_idx) + added_skips.add(next_idx) + result_indices = filled + return [self.tokens[i] for i in result_indices] + + def is_token_start(self, pos: int) -> bool: + """Return True if pos is the first character of a layout token (not mid-token).""" + if pos <= 0: + return True + if pos >= len(self._token_index_at): + return False + return self._token_index_at[pos - 1] != self._token_index_at[pos] + + +def _build_token_index(layout_document: LayoutDocument) -> _TokenIndex: + all_tokens: List[LayoutToken] = [] + skip_tokens: Set[int] = set() + no_space_after: Set[int] = set() + for line in layout_document.iter_all_lines(): + line_tokens: List[LayoutToken] = line.tokens or [] + for i, token in enumerate(line_tokens): + tok_global_idx = len(all_tokens) + all_tokens.append(token) + if i == len(line_tokens) - 1: + norm = normalize_for_alignment(token.text) + if norm == '-': + # Bare end-of-line hyphen produced by the PDF tokenizer when + # a word is split across lines (e.g. ["hyphen", "-"] then + # ["ation"]). Skip the "-" entirely and suppress the trailing + # space on the preceding word so the two halves join without a + # gap: "hyphen" + "ation" → "hyphenation". + skip_tokens.add(tok_global_idx) + if i > 0: + no_space_after.add(tok_global_idx - 1) + return _TokenIndex(all_tokens, skip_tokens=skip_tokens, no_space_after=no_space_after) + + +def _match_quality( + matching_blocks: List[Tuple[int, int, int]], + needle_len: int, +) -> float: + """Fraction of needle characters matched (0..1).""" + if needle_len == 0: + return 1.0 + matched = sum(size for _, _, size in matching_blocks if size) + return matched / needle_len + + +def _fuzzy_search_in_window( + haystack: str, + needle: str, + window_start: int, + window_end: int, + threshold: float, +) -> Optional[_MatchResult]: + """Try to find `needle` in haystack[window_start:window_end]. + + Returns (abs_start, abs_end, [(block_start, block_end), ...]) if quality >= + threshold, else None. Block ranges are in absolute haystack coordinates. + """ + window = haystack[window_start:window_end] + sm = LocalSequenceMatcher(a=window, b=needle, scoring=_SCORING) + blocks = sm.get_matching_blocks() + quality = _match_quality(blocks, len(needle)) + if quality < threshold: + return None + matched_blocks = [(ai, bi, size) for ai, bi, size in blocks if size] + if not matched_blocks: + return None + a_start = matched_blocks[0][0] + window_start + last = matched_blocks[-1] + a_end = last[0] + last[2] + window_start + abs_block_ranges: List[Tuple[int, int]] = [ + (ai + window_start, ai + size + window_start) + for ai, _bi, size in matched_blocks + ] + return a_start, a_end, abs_block_ranges + + +def _fuzzy_match_field_value( + token_index: _TokenIndex, + field_value: JatsFieldValue, + config: AlignmentConfig, + search_start: int, + search_end: Optional[int] = None, +) -> Optional[_MatchResult]: + needle = normalize_for_alignment(field_value.text) + if not needle: + return None + + haystack = token_index.haystack + hay_end = len(haystack) if search_end is None else min(search_end, len(haystack)) + + need_len = len(needle) + + window_size = max( + _DEFAULT_MIN_WINDOW, + min(config.max_window, need_len * _WINDOW_NEEDLE_MULTIPLIER), + ) + stride = max(1, window_size - need_len - 20) + + start = search_start + while start < hay_end: + end = min(start + window_size, hay_end) + result = _fuzzy_search_in_window(haystack, needle, start, end, config.threshold) + if result is not None: + return result + if end >= hay_end: + break + start += stride + + return None + + +def _search_range( + fv: JatsFieldValue, + last_match_end: int, + body_floor: int, + body_content_end: int, + front_matter_end: int, + keywords_floor: int, + reference_floor: int, + parent_match_by_field: Dict[str, Tuple[int, int]], +) -> Tuple[int, Optional[int]]: + """Return (search_start, search_end) for fv given current position state.""" + if fv.sub_field_name is not None and fv.field_name in parent_match_by_field: + p_start, p_end = parent_match_by_field[fv.field_name] + return p_start, p_end + _SUB_FIELD_PARENT_BUFFER + if fv.field_name in _BODY_CONTENT_FIELDS: + return max(0, max(body_floor, body_content_end) - 200), None + if fv.field_name in _REFERENCE_ANCHOR_FIELDS or fv.field_name in _REFERENCE_FIELDS: + # References use a dedicated floor that is independent of body_content_end. + # This prevents appendix or late body content from advancing body_content_end + # past the reference section, which would make the reference search start + # skip over all reference positions. + ref_start = max(0, reference_floor - 200) if reference_floor > 0 else max(0, body_floor) + return ref_start, None + if fv.field_name in _ANCHOR_FIELDS or fv.field_name in _POST_BODY_FIELDS: + # Anchor fields (abstract, title) and post-body fields (sub-articles) both + # search from last_match_end so they follow reading order and cannot fall + # back to the front-matter window. + return max(0, last_match_end - 200), None + if front_matter_end > 0: + # Front-matter constrained fields (authors, affs, keywords). + # Keywords are anchored to just after the keywords header/abstract so + # short common words ("confidence", "Bayesian") don't false-match in the + # abstract. Other front-matter fields start from position 0. + is_keywords = fv.field_name in _KEYWORDS_FIELDS + start = max(keywords_floor, front_matter_end) if is_keywords else 0 + return start, front_matter_end + _FRONT_MATTER_BUFFER + return max(0, last_match_end - 200), None + + +def _pre_anchor_indices( + block_ranges: List[Tuple[int, int]] +) -> Optional[Set[int]]: + """Return indices of non-anchor blocks tightly preceding the first anchor. + + Returns None when there are no anchor blocks at all (caller should label + every block unconditionally). Otherwise walks backward from the block + just before the first anchor; stops when the gap exceeds + _MAX_HAYSTACK_GAP_TO_FILL. The result covers DOI-prefix segments + ("10", ".", "1128", "/") that precede a longer segment but should still + be labeled. + """ + first_anchor_idx = next( + (i for i, (bs, be) in enumerate(block_ranges) if be - bs >= _MIN_ANCHOR_BLOCK_SIZE), + None, + ) + if first_anchor_idx is None: + return None + result: Set[int] = set() + prev_start = block_ranges[first_anchor_idx][0] + for i in range(first_anchor_idx - 1, -1, -1): + _, be = block_ranges[i] + if prev_start - be <= _MAX_HAYSTACK_GAP_TO_FILL: + result.add(i) + prev_start = block_ranges[i][0] + else: + break + return result + + +def _is_haystack_token_start(token_index: _TokenIndex, pos: int) -> bool: + """Return True if pos is the first character of a layout token in the haystack. + + A SW matching block may start mid-token when a single character in the tail of + a longer token happens to match the first character of the needle (e.g. the + final 'o' of "introdução" matching needle "o crescimento..."). Including such + a block in the pre-anchor fill would cause tokens_in_range to return the preceding + token and overwrite its label — typically a heading label with a paragraph label. + + DOI sub-tokens ("10", ".", "3233", "/") always occupy their own token and always + start at a token boundary, so they are unaffected by this guard. + """ + return token_index.is_token_start(pos) + + +def _label_tokens_for_blocks( + annotated: JatsAnnotatedLayoutDocument, + token_index: _TokenIndex, + block_ranges: List[Tuple[int, int]], + field_name: str, + sub_field_name: Optional[str], + instance_id: int, +) -> None: + """Label tokens using anchor+chain strategy (see module constants for rationale). + + The forward anchor+chain is extended by a backward pre-anchor pass: non-anchor + blocks that immediately precede the first anchor (tight gap ≤ _MAX_HAYSTACK_GAP_TO_FILL) + are included, so dense sub-tokens like DOI segments ("10", ".", "1128", "/") + that appear before the first long segment are not silently dropped. Blocks with + a larger gap before the first anchor (sidebar text, page headers) remain excluded. + + Pre-anchor blocks that start mid-token (tail characters of a longer token + incidentally matching the needle start) are skipped to prevent overwriting labels + already set by earlier field values on that token. + """ + pre_anchor = _pre_anchor_indices(block_ranges) + if pre_anchor is None: + # No anchor blocks at all: label everything so short field values are preserved. + for block_start, block_end in block_ranges: + for token in token_index.tokens_in_range(block_start, block_end): + annotated.set_token_label(token, field_name, sub_field_name, instance_id) + return + prev_included_end: Optional[int] = None + for i, (block_start, block_end) in enumerate(block_ranges): + is_anchor = (block_end - block_start) >= _MIN_ANCHOR_BLOCK_SIZE + within_gap = ( + prev_included_end is not None + and block_start - prev_included_end <= _MAX_HAYSTACK_GAP_TO_FILL + ) + if not is_anchor and not within_gap and i not in pre_anchor: + continue + if i in pre_anchor and not within_gap and not _is_haystack_token_start( + token_index, block_start + ): + continue + fill_start = prev_included_end if within_gap else block_start + assert fill_start is not None + for token in token_index.tokens_in_range(fill_start, block_end): + annotated.set_token_label(token, field_name, sub_field_name, instance_id) + prev_included_end = block_end + + +class LayoutDocumentJatsAligner: + """Aligns JATS field values to LayoutDocument tokens via fuzzy text matching.""" + + def __init__(self, config: Optional[AlignmentConfig] = None) -> None: + self.config = config or AlignmentConfig() + + def align( # pylint: disable=too-many-locals,too-many-branches + self, + layout_document: LayoutDocument, + field_values: List[JatsFieldValue], + ) -> JatsAnnotatedLayoutDocument: + annotated = JatsAnnotatedLayoutDocument(layout_document=layout_document) + if not field_values: + return annotated + + token_index = _build_token_index(layout_document) + if not token_index.haystack: + return annotated + + last_match_end = 0 + body_floor = 0 + body_content_end = 0 + front_matter_end = 0 + keywords_floor = 0 + reference_floor = 0 + parent_match_by_field: Dict[str, Tuple[int, int]] = {} + missed_by_field: Dict[str, int] = {} + matched_count = 0 + instance_by_field: Dict[str, int] = {} + + for fv in field_values: + search_start, search_end = _search_range( + fv, last_match_end, body_floor, body_content_end, + front_matter_end, keywords_floor, reference_floor, + parent_match_by_field, + ) + match_range = _fuzzy_match_field_value( + token_index, fv, self.config, + search_start=search_start, search_end=search_end, + ) + # Front-matter region constraint is soft: if a field value (e.g. an + # affiliation that appears near the end of the paper) is not found + # within the preferred region, fall back to a global search. Sub-field + # containment (search_end set because sub_field_name is not None) is a + # hard constraint and does not get this fallback. + if match_range is None and search_end is not None and fv.sub_field_name is None: + match_range = _fuzzy_match_field_value( + token_index, fv, self.config, search_start=0, search_end=None, + ) + # Body-content incremental constraint is also soft: body_content_end + # can jump forward when a nested sub-section (e.g. a mathematical + # appendix) matches at a later PDF position than subsequent paragraphs + # of the parent section. Fall back to searching from body_floor + # (end of abstract) so those paragraphs are not permanently blocked. + if ( + match_range is None + and fv.sub_field_name is None + and fv.field_name in _BODY_CONTENT_FIELDS + and search_start > body_floor + ): + match_range = _fuzzy_match_field_value( + token_index, fv, self.config, + search_start=body_floor, search_end=None, + ) + if match_range is None: + if fv.sub_field_name is None: + missed_by_field[fv.field_name] = ( + missed_by_field.get(fv.field_name, 0) + 1 + ) + continue + matched_count += 1 + a_start, a_end, block_ranges = match_range + last_match_end = max(last_match_end, a_end) + if fv.field_name in _ANCHOR_FIELDS: + body_floor = max(body_floor, a_end) + if fv.field_name in _FRONT_MATTER_END_FIELDS: + front_matter_end = max(front_matter_end, a_end) + if ( + fv.field_name in _KEYWORDS_SECTION_ANCHOR_FIELDS + or fv.field_name in _KEYWORDS_FIELDS + ): + keywords_floor = max(keywords_floor, a_end) + if fv.field_name in _BODY_CONTENT_FIELDS: + body_content_end = max(body_content_end, a_end) + if fv.field_name in _REFERENCE_ANCHOR_FIELDS or fv.field_name in _REFERENCE_FIELDS: + reference_floor = max(reference_floor, a_end) + if fv.sub_field_name is None: + parent_match_by_field[fv.field_name] = (a_start, a_end) + instance_by_field[fv.field_name] = ( + instance_by_field.get(fv.field_name, 0) + 1 + ) + instance_id = instance_by_field.get(fv.field_name, 0) + _label_tokens_for_blocks( + annotated, token_index, block_ranges, + fv.field_name, fv.sub_field_name, instance_id, + ) + + total = len(field_values) + if missed_by_field: + missed = sum(missed_by_field.values()) + LOGGER.warning( + 'Unmatched fields (%d/%d): %s', + missed, total, + ', '.join('%s:%d' % (k, v) for k, v in sorted(missed_by_field.items())), + ) + else: + LOGGER.info('Aligned all %d field values', total) + + return annotated diff --git a/sciencebeam_parser/training/jats/annotated_document.py b/sciencebeam_parser/training/jats/annotated_document.py new file mode 100644 index 00000000..e6680cbf --- /dev/null +++ b/sciencebeam_parser/training/jats/annotated_document.py @@ -0,0 +1,41 @@ +from dataclasses import dataclass, field +from typing import Dict, Optional, Tuple + +from sciencebeam_parser.document.layout_document import LayoutDocument, LayoutToken + + +# id(token) -> (field_name, sub_field_name_or_None, instance_id) +TokenLabelById = Dict[int, Tuple[str, Optional[str], int]] + + +@dataclass +class JatsAnnotatedLayoutDocument: + layout_document: LayoutDocument + token_label_by_id: TokenLabelById = field(default_factory=dict) + + def get_token_field(self, token: LayoutToken) -> Optional[str]: + entry = self.token_label_by_id.get(id(token)) + return entry[0] if entry is not None else None + + def get_token_sub_field(self, token: LayoutToken) -> Optional[str]: + entry = self.token_label_by_id.get(id(token)) + return entry[1] if entry is not None else None + + def get_token_instance(self, token: LayoutToken) -> int: + entry = self.token_label_by_id.get(id(token)) + return entry[2] if entry is not None else 0 + + def set_token_label( + self, + token: LayoutToken, + field_name: str, + sub_field_name: Optional[str] = None, + instance_id: int = 0, + ) -> None: + self.token_label_by_id[id(token)] = (field_name, sub_field_name, instance_id) + + def coverage_ratio(self) -> float: + total = sum(1 for _ in self.layout_document.iter_all_tokens()) + if total == 0: + return 1.0 + return len(self.token_label_by_id) / total diff --git a/sciencebeam_parser/training/jats/coverage.py b/sciencebeam_parser/training/jats/coverage.py new file mode 100644 index 00000000..f453fae7 --- /dev/null +++ b/sciencebeam_parser/training/jats/coverage.py @@ -0,0 +1,70 @@ +from dataclasses import dataclass, field +from typing import Collection, Mapping, Set + +from sciencebeam_parser.training.jats.annotated_document import JatsAnnotatedLayoutDocument + + +@dataclass +class CoverageResult: + """Summary of how well a set of required fields was matched.""" + required_fields_present: Set[str] = field(default_factory=set) + required_fields_missing: Set[str] = field(default_factory=set) + required_matching_fields_matched: Set[str] = field(default_factory=set) + required_matching_fields_missing: Set[str] = field(default_factory=set) + + @property + def is_passing(self) -> bool: + return ( + not self.required_fields_missing + and not self.required_matching_fields_missing + ) + + def __str__(self) -> str: + parts = [] + if self.required_fields_missing: + parts.append(f'required fields absent: {sorted(self.required_fields_missing)}') + if self.required_matching_fields_missing: + parts.append( + f'matching fields not aligned: ' + f'{sorted(self.required_matching_fields_missing)}' + ) + return '; '.join(parts) if parts else 'OK' + + +def check_coverage( + annotated: JatsAnnotatedLayoutDocument, + field_values_by_field: Mapping[str, bool], + required_fields: Collection[str], + require_matching_fields: Collection[str], +) -> CoverageResult: + """ + Args: + annotated: the annotated layout document + field_values_by_field: mapping of field_name → whether that field appeared in JATS + required_fields: fields that must be present AND aligned + require_matching_fields: fields that must be aligned IF they appear in JATS + """ + aligned_fields: Set[str] = { + entry[0] + for entry in annotated.token_label_by_id.values() + } + present_fields: Set[str] = { + f for f, present in field_values_by_field.items() if present + } + + result = CoverageResult() + for f in required_fields: + if f not in present_fields or f not in aligned_fields: + result.required_fields_missing.add(f) + else: + result.required_fields_present.add(f) + + for f in require_matching_fields: + if f not in present_fields: + continue # not in JATS → no constraint + if f not in aligned_fields: + result.required_matching_fields_missing.add(f) + else: + result.required_matching_fields_matched.add(f) + + return result diff --git a/sciencebeam_parser/training/jats/field_extractor.py b/sciencebeam_parser/training/jats/field_extractor.py new file mode 100644 index 00000000..bb0eef84 --- /dev/null +++ b/sciencebeam_parser/training/jats/field_extractor.py @@ -0,0 +1,404 @@ +from typing import Dict, Iterator, List, Optional, Sequence, Tuple +from dataclasses import dataclass + +from lxml import etree + +from sciencebeam_parser.training.jats.field_vocab import ( + JatsFieldNames, + JatsSubFieldNames, +) + + +@dataclass +class JatsFieldValue: + text: str + field_name: str + sub_field_name: Optional[str] = None + + +def _element_text(el: etree._Element) -> str: + return ' '.join(' '.join(el.itertext()).split()) + + +def _iter_sub_field_values( + parent_el: etree._Element, + field_name: str, + sub_xpath_by_sub_field: Sequence[tuple], +) -> Iterator[JatsFieldValue]: + """Yield one JatsFieldValue per sub-field span found in parent_el.""" + for sub_field_name, xpath in sub_xpath_by_sub_field: + for child in parent_el.xpath(xpath): + text = _element_text(child) + if text: + yield JatsFieldValue( + text=text, + field_name=field_name, + sub_field_name=sub_field_name, + ) + + +# Sub-field XPaths for references (relative to each element) +_REFERENCE_SUB_FIELDS = [ + (JatsSubFieldNames.REFERENCE_LABEL, './label'), + (JatsSubFieldNames.REFERENCE_AUTHOR, './/string-name[not(ancestor::person-group)]'), + (JatsSubFieldNames.REFERENCE_ARTICLE_TITLE, './/article-title'), + (JatsSubFieldNames.REFERENCE_SOURCE, './/source'), + (JatsSubFieldNames.REFERENCE_YEAR, './/year'), + (JatsSubFieldNames.REFERENCE_VOLUME, './/volume'), + (JatsSubFieldNames.REFERENCE_ISSUE, './/issue'), + (JatsSubFieldNames.REFERENCE_FPAGE, './/fpage'), + (JatsSubFieldNames.REFERENCE_LPAGE, './/lpage'), + (JatsSubFieldNames.REFERENCE_PUBLISHER_NAME, './/publisher-name'), + (JatsSubFieldNames.REFERENCE_PUBLISHER_LOC, './/publisher-loc'), + (JatsSubFieldNames.REFERENCE_DOI, './/pub-id[@pub-id-type="doi"]'), + (JatsSubFieldNames.REFERENCE_PMID, './/pub-id[@pub-id-type="pmid"]'), + (JatsSubFieldNames.REFERENCE_PMCID, './/pub-id[@pub-id-type="pmcid"]'), + (JatsSubFieldNames.REFERENCE_WEB, + './/ext-link[@ext-link-type="uri"][starts-with(normalize-space(.), "http")]'), +] + +# Sub-field XPaths for affiliations (relative to each aff element) +_AFF_SUB_FIELDS = [ + (JatsSubFieldNames.AUTHOR_AFF_LABEL, './label'), + (JatsSubFieldNames.AUTHOR_AFF_INSTITUTION, './institution'), + (JatsSubFieldNames.AUTHOR_AFF_DEPARTMENT, + './addr-line/named-content[@content-type="department"]'), + (JatsSubFieldNames.AUTHOR_AFF_CITY, + './addr-line/named-content[@content-type="city"]'), + (JatsSubFieldNames.AUTHOR_AFF_POSTCODE, + './addr-line/named-content[@content-type="postcode"]'), + (JatsSubFieldNames.AUTHOR_AFF_REGION, './addr-line/named-content[@content-type="state"]'), + (JatsSubFieldNames.AUTHOR_AFF_COUNTRY, './country'), +] + + +def _local_tag(el: etree._Element) -> str: + tag = el.tag + if isinstance(tag, str) and tag.startswith('{'): + return tag.split('}', 1)[1] + return tag if isinstance(tag, str) else '' + + +def _aff_addr_parts(aff_el: etree._Element) -> List[str]: + """Collect address text from an in document order. + + Covers three JATS patterns: + - Structured: and/or elements + - Semi-structured: present but city/postcode sit in its tail text + (no ), e.g. 'UCL, London WC1N 1EH, + UK' + - Unstructured (label-only affs): returns nothing; address cannot be determined + + Institution tail text is only included when the aff also has a or + element. Without that anchor the tail may be continuation of the + institution name rather than a geographic address (e.g. a department name split + across two tags). + """ + has_structured_addr = bool(aff_el.xpath('./country | ./addr-line')) + parts: List[str] = [] + for child in aff_el: + tag = _local_tag(child) + if tag in ('addr-line', 'country'): + text = _element_text(child) + if text: + parts.append(text) + elif tag == 'institution' and has_structured_addr: + # Tail text after is city/postcode when no is present. + # Only collected when a or confirms this aff has structured + # address content, to avoid misclassifying department-name continuations. + tail = ' '.join((child.tail or '').split()).strip(', ') + if tail: + parts.append(tail) + return parts + + +def _iter_aff_elements(root: etree._Element) -> Iterator[etree._Element]: + yield from root.xpath( + 'front/article-meta/contrib-group/aff' + '| front/article-meta/contrib-group/contrib/aff' + '| front/article-meta/aff' + ) + + +class JatsFieldExtractor: + """Extract (text, field_name, sub_field_name) triples from a JATS
root.""" + + def iter_field_values(self, root: etree._Element) -> Iterator[JatsFieldValue]: + yield from self._iter_front_values(root) + yield from self._iter_body_values(root) + yield from self._iter_back_values(root) + yield from self._iter_sub_article_values(root) + + def _emit( + self, + elements: List[etree._Element], + field_name: str, + ) -> Iterator[JatsFieldValue]: + for el in elements: + text = _element_text(el) + if text: + yield JatsFieldValue(text=text, field_name=field_name) + + # ── Front matter ────────────────────────────────────────────────────────── + + def _iter_front_values(self, root: etree._Element) -> Iterator[JatsFieldValue]: + yield from self._iter_front_metadata_values(root) + yield from self._iter_front_contrib_values(root) + + def _iter_front_metadata_values(self, root: etree._Element) -> Iterator[JatsFieldValue]: + for el in root.xpath('front/article-meta/title-group/article-title'): + text = _element_text(el) + if text: + yield JatsFieldValue(text=text, field_name=JatsFieldNames.TITLE) + + for el in root.xpath('front/article-meta/abstract'): + text = _element_text(el) + if text: + yield JatsFieldValue(text=text, field_name=JatsFieldNames.ABSTRACT) + + # Per GROBID annotation guidelines, the "Keywords" heading is not annotated + # in the header model. It is still emitted as KEYWORDS_TITLE so that the + # segmentation model can label it as
. Combine all children + # of a into a single KEYWORDS field value. + for kwd_group in root.xpath('front/article-meta/kwd-group'): + for title_el in kwd_group.xpath('./title'): + text = _element_text(title_el) + if text: + yield JatsFieldValue(text=text, field_name=JatsFieldNames.KEYWORDS_TITLE) + kwd_texts = [ + _element_text(kwd_el) + for kwd_el in kwd_group.xpath('./kwd') + if _element_text(kwd_el) + ] + if kwd_texts: + yield JatsFieldValue( + text=', '.join(kwd_texts), + field_name=JatsFieldNames.KEYWORDS, + ) + + for el in root.xpath( + 'front/article-meta/article-categories' + '/subj-group/subject[@subj-group-type="display-channel"]' + ): + text = _element_text(el) + if text: + yield JatsFieldValue(text=text, field_name=JatsFieldNames.MANUSCRIPT_TYPE) + + yield from self._iter_front_publication_values(root) + + def _iter_front_publication_values(self, root: etree._Element) -> Iterator[JatsFieldValue]: + """Funding statements and copyright / licence text from front matter.""" + for el in root.xpath('front/article-meta/funding-group/funding-statement'): + text = _element_text(el) + if text: + yield JatsFieldValue(text=text, field_name=JatsFieldNames.FUNDING) + + for el in root.xpath( + 'front/article-meta/permissions/copyright-statement' + ' | front/article-meta/permissions/license/license-p' + ): + text = _element_text(el) + if text: + yield JatsFieldValue(text=text, field_name=JatsFieldNames.COPYRIGHT) + + def _iter_front_contrib_values(self, root: etree._Element) -> Iterator[JatsFieldValue]: + # Per GROBID annotation guidelines, all author tokens in the byline (including + # affiliation markers and separating punctuation) are labelled . + # Emit one merged field value per contrib-group so the aligner covers the + # whole byline span, including commas, "&", etc. between individual names. + # JATS stores names in Surname-Given order; PDFs display Given-Surname, so + # each name part is reversed. Affiliation/fn/corresp xref markers are + # appended to each name so the combined needle matches the PDF author line. + for contrib_group in root.xpath('front/article-meta/contrib-group'): + author_parts = [] + for contrib in contrib_group.xpath( + 'contrib[not(@contrib-type) or @contrib-type="author"]' + ): + name_el = contrib.find('name') + if name_el is None: + continue + given = (name_el.findtext('given-names') or '').strip() + surname = (name_el.findtext('surname') or '').strip() + name_text = ( + ' '.join(p for p in [given, surname] if p) or _element_text(name_el) + ) + if not name_text: + continue + markers = [ + x.text.strip() + for x in contrib.xpath( + 'xref[@ref-type="aff" or @ref-type="fn" or @ref-type="corresp"]' + ) + if x.text and x.text.strip() + ] + author_parts.append(' '.join([name_text] + markers)) + if author_parts: + yield JatsFieldValue( + text=' '.join(author_parts), + field_name=JatsFieldNames.AUTHOR, + ) + + for aff_el in _iter_aff_elements(root): + text = _element_text(aff_el) + if text: + yield JatsFieldValue(text=text, field_name=JatsFieldNames.AUTHOR_AFF) + # Emit a bulk address value BEFORE individual sub-fields so that commas + # between city and country also get the AUTHOR_AFF_ADDR sub-field label + # (individual city/country sub-fields overwrite their own tokens afterward). + addr_texts = _aff_addr_parts(aff_el) + if addr_texts: + yield JatsFieldValue( + text=' '.join(addr_texts), + field_name=JatsFieldNames.AUTHOR_AFF, + sub_field_name=JatsSubFieldNames.AUTHOR_AFF_ADDR, + ) + yield from _iter_sub_field_values( + aff_el, JatsFieldNames.AUTHOR_AFF, _AFF_SUB_FIELDS + ) + + for el in root.xpath('front/article-meta/author-notes/*'): + text = _element_text(el) + if text: + yield JatsFieldValue(text=text, field_name=JatsFieldNames.AUTHOR_NOTES) + + for el in root.xpath('front/article-meta/fpage | front/article-meta/lpage'): + text = _element_text(el) + if text: + yield JatsFieldValue(text=text, field_name=JatsFieldNames.PAGE_NO) + + # ── Body ────────────────────────────────────────────────────────────────── + + def _iter_body_values(self, root: etree._Element) -> Iterator[JatsFieldValue]: + body = root.find('body') + if body is None: + return + # Build document-order index so that section titles, paragraphs, figures, + # and tables are yielded interleaved as they appear in the XML, not grouped + # by type. If section titles are all emitted first the aligner's + # body_content_end advances past the early paragraphs before they are matched. + position: Dict[etree._Element, int] = {el: i for i, el in enumerate(root.iter())} + entries: List[Tuple[int, JatsFieldValue]] = [] + + for el in body.xpath('.//sec/title'): + text = _element_text(el) + if text: + entries.append((position[el], JatsFieldValue( + text=text, field_name=JatsFieldNames.BODY_SECTION_TITLE))) + + for el in body.xpath('.//p[not(ancestor::fig) and not(ancestor::table-wrap)]'): + text = _element_text(el) + if text: + entries.append((position[el], JatsFieldValue( + text=text, field_name=JatsFieldNames.BODY_SECTION_PARAGRAPH))) + + for el in body.xpath('.//fig'): + children = el.xpath('./label') + el.xpath('./caption') + text = (_element_text(el) if not children + else ' '.join(_element_text(c) for c in children if _element_text(c))) + if text: + entries.append((position[el], JatsFieldValue( + text=text, field_name=JatsFieldNames.BODY_FIGURE))) + + for el in body.xpath('.//table-wrap'): + children = el.xpath('./label') + el.xpath('./caption') + text = (_element_text(el) if not children + else ' '.join(_element_text(c) for c in children if _element_text(c))) + if text: + entries.append((position[el], JatsFieldValue( + text=text, field_name=JatsFieldNames.BODY_TABLE))) + + for _, fv in sorted(entries): + yield fv + + # ── Back matter ─────────────────────────────────────────────────────────── + + def _iter_back_values(self, root: etree._Element) -> Iterator[JatsFieldValue]: + yield from self._iter_back_narrative_values(root) + yield from self._iter_back_reference_values(root) + + def _iter_back_narrative_values( # pylint: disable=too-many-branches + self, root: etree._Element + ) -> Iterator[JatsFieldValue]: + position: Dict[etree._Element, int] = {el: i for i, el in enumerate(root.iter())} + entries: List[Tuple[int, JatsFieldValue]] = [] + + for el in root.xpath('//ack//title'): + text = _element_text(el) + if text: + entries.append((position[el], JatsFieldValue( + text=text, field_name=JatsFieldNames.ACK_SECTION_TITLE))) + + for el in root.xpath('//ack//p'): + text = _element_text(el) + if text: + entries.append((position[el], JatsFieldValue( + text=text, field_name=JatsFieldNames.ACK_SECTION_PARAGRAPH))) + + for el in root.xpath('//app-group/title'): + text = _element_text(el) + if text: + entries.append((position[el], JatsFieldValue( + text=text, field_name=JatsFieldNames.APPENDIX_GROUP_TITLE))) + + for el in root.xpath('//app'): + text = _element_text(el) + if text: + entries.append((position[el], JatsFieldValue( + text=text, field_name=JatsFieldNames.APPENDIX))) + + for el in root.xpath('back//sec[not(ancestor::ack)]/title'): + text = _element_text(el) + if text: + entries.append((position[el], JatsFieldValue( + text=text, field_name=JatsFieldNames.BACK_SECTION_TITLE))) + + for el in root.xpath( + 'back//sec[not(ancestor::ack)]/p[not(ancestor::ack)]' + ' | back//p[not(ancestor::sec) and not(ancestor::ack)]' + ): + text = _element_text(el) + if text: + entries.append((position[el], JatsFieldValue( + text=text, field_name=JatsFieldNames.BACK_SECTION_PARAGRAPH))) + + for _, fv in sorted(entries): + yield fv + + def _iter_back_reference_values(self, root: etree._Element) -> Iterator[JatsFieldValue]: + for el in root.xpath('back/ref-list/title'): + text = _element_text(el) + if text: + yield JatsFieldValue( + text=text, field_name=JatsFieldNames.REFERENCE_LIST_TITLE + ) + + for ref_el in root.xpath('back/ref-list/ref'): + text = _element_text(ref_el) + if text: + yield JatsFieldValue(text=text, field_name=JatsFieldNames.REFERENCE) + yield from _iter_sub_field_values( + ref_el, JatsFieldNames.REFERENCE, _REFERENCE_SUB_FIELDS + ) + + # ── Sub-articles (ORE peer-review reports, etc.) ────────────────────────── + + def _iter_sub_article_values(self, root: etree._Element) -> Iterator[JatsFieldValue]: + """Yield paragraph/title text from elements as SUB_ARTICLE values. + + ORE papers embed peer-review reports as sub-articles. Extracting their + content in document order lets the aligner map those PDF pages to the + SUB_ARTICLE field, which the segmentation model labels as rather + than . + """ + position: Dict[etree._Element, int] = {el: i for i, el in enumerate(root.iter())} + entries: List[Tuple[int, JatsFieldValue]] = [] + + for sub_article in root.xpath('.//sub-article'): + for el in sub_article.xpath('.//title | .//p'): + text = _element_text(el) + if text: + entries.append((position[el], JatsFieldValue( + text=text, field_name=JatsFieldNames.SUB_ARTICLE))) + + for _, fv in sorted(entries): + yield fv diff --git a/sciencebeam_parser/training/jats/field_vocab.py b/sciencebeam_parser/training/jats/field_vocab.py new file mode 100644 index 00000000..a5c6761e --- /dev/null +++ b/sciencebeam_parser/training/jats/field_vocab.py @@ -0,0 +1,137 @@ +from typing import Dict + + +class JatsFieldNames: + TITLE = 'title' + ABSTRACT = 'abstract' + KEYWORDS_TITLE = 'keywords_title' + KEYWORDS = 'keywords' + MANUSCRIPT_TYPE = 'manuscript_type' + AUTHOR = 'author' + AUTHOR_AFF = 'author_aff' + AUTHOR_NOTES = 'author_notes' + FUNDING = 'funding' + COPYRIGHT = 'copyright' + SUB_ARTICLE = 'sub_article' + BODY_SECTION_TITLE = 'body_section_title' + BODY_SECTION_PARAGRAPH = 'body_section_paragraph' + BODY_FIGURE = 'body_figure' + BODY_TABLE = 'body_table' + BACK_SECTION_TITLE = 'back_section_title' + BACK_SECTION_PARAGRAPH = 'back_section_paragraph' + ACK_SECTION_TITLE = 'acknowledgment_section_title' + ACK_SECTION_PARAGRAPH = 'acknowledgment_section_paragraph' + APPENDIX_GROUP_TITLE = 'appendix_group_title' + APPENDIX = 'appendix' + REFERENCE_LIST_TITLE = 'reference_list_title' + REFERENCE = 'reference' + PAGE_NO = 'page_no' + + +class JatsSubFieldNames: + AUTHOR_AFF_ADDR = 'author_aff_addr' + REFERENCE_AUTHOR = 'reference-author' + REFERENCE_ARTICLE_TITLE = 'reference-article-title' + REFERENCE_SOURCE = 'reference-source' + REFERENCE_YEAR = 'reference-year' + REFERENCE_VOLUME = 'reference-volume' + REFERENCE_ISSUE = 'reference-issue' + REFERENCE_FPAGE = 'reference-fpage' + REFERENCE_LPAGE = 'reference-lpage' + REFERENCE_DOI = 'reference-doi' + REFERENCE_PMID = 'reference-pmid' + REFERENCE_PMCID = 'reference-pmcid' + REFERENCE_WEB = 'reference-web' + REFERENCE_LABEL = 'reference-label' + REFERENCE_PUBLISHER_NAME = 'reference-publisher-name' + REFERENCE_PUBLISHER_LOC = 'reference-publisher-loc' + AUTHOR_AFF_LABEL = 'author_aff-label' + AUTHOR_AFF_INSTITUTION = 'author_aff-institution' + AUTHOR_AFF_DEPARTMENT = 'author_aff-department' + AUTHOR_AFF_CITY = 'author_aff-address-city' + AUTHOR_AFF_POSTCODE = 'author_aff-address-postcode' + AUTHOR_AFF_REGION = 'author_aff-address-state' + AUTHOR_AFF_COUNTRY = 'author_aff-address-country' + + +# ── Segmentation label mapping (mirrors segmentation.conf [tags]) ───────────── +SEGMENTATION_LABEL_BY_FIELD: Dict[str, str] = { + JatsFieldNames.TITLE: '
', + JatsFieldNames.ABSTRACT: '
', + JatsFieldNames.KEYWORDS_TITLE: '
', + JatsFieldNames.KEYWORDS: '
', + JatsFieldNames.MANUSCRIPT_TYPE: '
', + JatsFieldNames.AUTHOR: '
', + JatsFieldNames.AUTHOR_AFF: '
', + JatsFieldNames.AUTHOR_NOTES: '
', + JatsFieldNames.FUNDING: '
', + JatsFieldNames.COPYRIGHT: '
', + JatsFieldNames.SUB_ARTICLE: '', + JatsFieldNames.BODY_SECTION_TITLE: '', + JatsFieldNames.BODY_SECTION_PARAGRAPH: '', + JatsFieldNames.BODY_FIGURE: '', + JatsFieldNames.BODY_TABLE: '', + JatsFieldNames.ACK_SECTION_TITLE: '', + JatsFieldNames.ACK_SECTION_PARAGRAPH: '', + JatsFieldNames.APPENDIX_GROUP_TITLE: '', + JatsFieldNames.APPENDIX: '', + JatsFieldNames.BACK_SECTION_TITLE: '', + JatsFieldNames.BACK_SECTION_PARAGRAPH: '', + JatsFieldNames.REFERENCE_LIST_TITLE: '', + JatsFieldNames.REFERENCE: '', + JatsFieldNames.PAGE_NO: '', +} + +# ── Header model label mapping ──────────────────────────────────────────────── +HEADER_LABEL_BY_FIELD: Dict[str, str] = { + JatsFieldNames.TITLE: '', + JatsFieldNames.ABSTRACT: '<abstract>', + JatsFieldNames.KEYWORDS: '<keyword>', + JatsFieldNames.AUTHOR: '<author>', + JatsFieldNames.AUTHOR_AFF: '<affiliation>', + JatsFieldNames.AUTHOR_NOTES: '<note>', + JatsFieldNames.MANUSCRIPT_TYPE: '<note>', +} + +# ── Fulltext model label mapping ────────────────────────────────────────────── +FULLTEXT_LABEL_BY_FIELD: Dict[str, str] = { + JatsFieldNames.BODY_SECTION_TITLE: '<section>', + JatsFieldNames.BODY_SECTION_PARAGRAPH: '<paragraph>', + JatsFieldNames.BODY_FIGURE: '<figure>', + JatsFieldNames.BODY_TABLE: '<table>', + JatsFieldNames.ACK_SECTION_TITLE: '<section>', + JatsFieldNames.ACK_SECTION_PARAGRAPH: '<paragraph>', + JatsFieldNames.BACK_SECTION_TITLE: '<section>', + JatsFieldNames.BACK_SECTION_PARAGRAPH: '<paragraph>', +} + +# ── Citation model label mapping (keyed by sub-field name) ──────────────────── +# Tokens whose sub_field_name matches get this label; all others → <note>. +CITATION_LABEL_BY_SUB_FIELD: Dict[str, str] = { + JatsSubFieldNames.REFERENCE_AUTHOR: '<author>', + JatsSubFieldNames.REFERENCE_ARTICLE_TITLE: '<title>', + JatsSubFieldNames.REFERENCE_SOURCE: '<journal>', + JatsSubFieldNames.REFERENCE_YEAR: '<date>', + JatsSubFieldNames.REFERENCE_VOLUME: '<volume>', + JatsSubFieldNames.REFERENCE_ISSUE: '<issue>', + JatsSubFieldNames.REFERENCE_FPAGE: '<pages>', + JatsSubFieldNames.REFERENCE_LPAGE: '<pages>', + JatsSubFieldNames.REFERENCE_DOI: '<idno>', + JatsSubFieldNames.REFERENCE_PMID: '<pubnum>', + JatsSubFieldNames.REFERENCE_PMCID: '<pubnum>', + JatsSubFieldNames.REFERENCE_WEB: '<web>', + JatsSubFieldNames.REFERENCE_LABEL: '<note>', + JatsSubFieldNames.REFERENCE_PUBLISHER_NAME: '<publisher>', + JatsSubFieldNames.REFERENCE_PUBLISHER_LOC: '<location>', +} + +# ── Affiliation-address model label mapping (keyed by sub-field name) ───────── +AFF_LABEL_BY_SUB_FIELD: Dict[str, str] = { + JatsSubFieldNames.AUTHOR_AFF_LABEL: '<marker>', + JatsSubFieldNames.AUTHOR_AFF_INSTITUTION: '<institution>', + JatsSubFieldNames.AUTHOR_AFF_DEPARTMENT: '<department>', + JatsSubFieldNames.AUTHOR_AFF_CITY: '<settlement>', + JatsSubFieldNames.AUTHOR_AFF_POSTCODE: '<postCode>', + JatsSubFieldNames.AUTHOR_AFF_REGION: '<region>', + JatsSubFieldNames.AUTHOR_AFF_COUNTRY: '<country>', +} diff --git a/sciencebeam_parser/training/jats/segmentation.py b/sciencebeam_parser/training/jats/segmentation.py new file mode 100644 index 00000000..07c2d1bf --- /dev/null +++ b/sciencebeam_parser/training/jats/segmentation.py @@ -0,0 +1,322 @@ +import logging +import re +from collections import Counter +from dataclasses import dataclass +from typing import Dict, List, Mapping, Optional, Set + +from sciencebeam_parser.document.layout_document import ( + LayoutDocument, + LayoutLine, + LayoutPageMeta, + LayoutToken, +) +from sciencebeam_parser.training.jats.annotated_document import JatsAnnotatedLayoutDocument +from sciencebeam_parser.training.jats.field_vocab import SEGMENTATION_LABEL_BY_FIELD + + +LOGGER = logging.getLogger(__name__) + +# Segmentation label constants (mirror SegmentationTagNames in trainer-grobid-tools) +SEG_FRONT = '<header>' +SEG_BODY = '<body>' +SEG_REFERENCES = '<references>' +SEG_ACKNOWLEDGEMENT = '<acknowledgement>' +SEG_ANNEX = '<annex>' +SEG_PAGE = '<page>' +SEG_HEADNOTE = '<headnote>' +SEG_FOOTNOTE = '<footnote>' + +# Fraction of page height: lines above this → headnote, below this → footnote candidate +_HEADNOTE_Y_RATIO = 0.08 +_FOOTNOTE_Y_RATIO = 0.92 + +# Line index threshold: front blocks starting beyond this are cleared. +# ORE papers have a second front-matter page (author roles, competing interests, +# grant info, copyright) that can start at line ~60+, so the threshold is set high +# enough to preserve those blocks when they match JATS front-matter fields. +_DEFAULT_FRONT_MAX_START_LINE_INDEX = 80 +# Headnotes are expected in the first few lines (index ≤ this) +_DEFAULT_PAGE_HEADER_MAX_FIRST_LINE_INDEX = 5 + + +@dataclass +class SegmentationConfig: + front_max_start_line_index: int = _DEFAULT_FRONT_MAX_START_LINE_INDEX + page_header_max_first_line_index: int = _DEFAULT_PAGE_HEADER_MAX_FIRST_LINE_INDEX + headnote_y_ratio: float = _HEADNOTE_Y_RATIO + footnote_y_ratio: float = _FOOTNOTE_Y_RATIO + + +@dataclass +class _SegLine: + layout_line: LayoutLine + line_index: int + seg_label: Optional[str] = None + + @property + def text(self) -> str: + return self.layout_line.text + + @property + def first_token(self) -> Optional[LayoutToken]: + tokens = self.layout_line.tokens + return tokens[0] if tokens else None + + +def _majority_vote_label( + tokens: List[LayoutToken], + annotated: JatsAnnotatedLayoutDocument, +) -> Optional[str]: + field_names: List[str] = [ + label + for t in tokens + if (label := annotated.get_token_field(t)) is not None + ] + if not field_names: + return None + most_common_field: str = Counter(field_names).most_common(1)[0][0] + return SEGMENTATION_LABEL_BY_FIELD.get(most_common_field) + + +def _is_valid_page_number_candidate(text: str) -> bool: + stripped = text.strip() + if not stripped: + return False + try: + int(stripped) + return True + except ValueError: + return False + + +def _parse_page_number(text: str) -> Optional[int]: + try: + return int(text.strip()) + except ValueError: + return None + + +def _is_valid_headnote_candidate(text: str, count: int, min_count: int = 2) -> bool: + if count < min_count: + return False + if re.match(r'^(\d|\s|\.)+$', text): + return False + if len(re.split(r'\s', text.strip())) < 2: + return False + return True + + +def _get_page_meta_by_page_number( + layout_document: LayoutDocument, +) -> Mapping[int, LayoutPageMeta]: + return { + page.meta.page_number: page.meta + for page in layout_document.pages + } + + +def _get_line_y_ratio( + seg_line: _SegLine, + page_meta_by_number: Mapping[int, LayoutPageMeta], +) -> Optional[float]: + token = seg_line.first_token + if token is None or token.coordinates is None or not token.coordinates: + return None + coords = token.coordinates + page_meta = page_meta_by_number.get(coords.page_number) + if page_meta is None or page_meta.coordinates is None or not page_meta.coordinates: + return None + page_height = page_meta.coordinates.height + if page_height <= 0: + return None + return coords.y / page_height + + +# ── Heuristic passes ────────────────────────────────────────────────────────── + +def _tag_by_coordinates( + seg_lines: List[_SegLine], + page_meta_by_number: Mapping[int, LayoutPageMeta], + config: SegmentationConfig, +) -> None: + """Use vertical position to label headnotes and footnotes for untagged lines.""" + for seg_line in seg_lines: + if seg_line.seg_label is not None: + continue + y_ratio = _get_line_y_ratio(seg_line, page_meta_by_number) + if y_ratio is None: + continue + if y_ratio < config.headnote_y_ratio: + seg_line.seg_label = SEG_HEADNOTE + elif y_ratio > config.footnote_y_ratio: + if _is_valid_page_number_candidate(seg_line.text): + seg_line.seg_label = SEG_PAGE + else: + seg_line.seg_label = SEG_FOOTNOTE + + +def _tag_headnotes_by_text_repetition( + seg_lines: List[_SegLine], + max_first_line_index: int, +) -> None: + """Text repetition fallback for headnote detection (no coordinates).""" + untagged_text_counts: Counter = Counter( + sl.text for sl in seg_lines if sl.seg_label is None + ) + if not untagged_text_counts: + return + min_count: Optional[int] = None + for text, count in untagged_text_counts.most_common(): + if not _is_valid_headnote_candidate(text, count, min_count=min_count or 2): + continue + first_line_index = next( + (sl.line_index for sl in seg_lines if sl.text == text), -1 + ) + if first_line_index >= max_first_line_index: + continue + if min_count is None: + min_count = max(2, count - 1) + for sl in seg_lines: + if sl.text == text and sl.seg_label is None: + sl.seg_label = SEG_HEADNOTE + + +def _find_missing_page_numbers(seg_lines: List[_SegLine]) -> None: + """Label standalone numeric untagged lines that fit between known page-number lines.""" + + @dataclass + class _Candidate: + seg_line: _SegLine + page_number: int + + existing = [ + _Candidate(sl, _parse_page_number(sl.text)) # type: ignore[arg-type] + for sl in seg_lines + if sl.seg_label == SEG_PAGE and _parse_page_number(sl.text) is not None + ] + candidates = [ + _Candidate(sl, _parse_page_number(sl.text)) # type: ignore[arg-type] + for sl in seg_lines + if sl.seg_label is None and _is_valid_page_number_candidate(sl.text) + ] + if not existing or not candidates: + return + + min_page_number = 1 + for known in existing: + max_line_index = known.seg_line.line_index + max_page_number = known.page_number - 1 + for cand in candidates: + if cand.seg_line.line_index >= max_line_index: + continue + if cand.page_number < min_page_number or cand.page_number > max_page_number: + continue + cand.seg_line.seg_label = SEG_PAGE + min_page_number = known.page_number + 1 + + +def _clear_front_beyond_threshold( + seg_lines: List[_SegLine], + max_block_start_line_index: int, +) -> None: + if not max_block_start_line_index: + return + block_label: Optional[str] = None + block_start_idx = 0 + for sl in seg_lines: + if sl.seg_label != block_label: + block_label = sl.seg_label + block_start_idx = sl.line_index + if ( + block_label == SEG_FRONT + and block_start_idx > max_block_start_line_index + ): + sl.seg_label = None + + +def _merge_gap_lines( + seg_lines: List[_SegLine], + enabled_labels: Set[str], + enabled_tail_labels: Set[str], +) -> None: + """Assign untagged gap lines to the surrounding region.""" + _IGNORED = {SEG_HEADNOTE, SEG_PAGE} + candidate_gap: List[_SegLine] = [] + prev_label: Optional[str] = SEG_FRONT + for sl in seg_lines: + if sl.seg_label in _IGNORED: + continue + if sl.seg_label is not None: + if prev_label == sl.seg_label and sl.seg_label in enabled_labels: + for gap_sl in candidate_gap: + gap_sl.seg_label = sl.seg_label + candidate_gap = [] + prev_label = sl.seg_label + elif prev_label in enabled_labels: + candidate_gap.append(sl) + else: + candidate_gap = [] + + if candidate_gap and prev_label in enabled_tail_labels: + for gap_sl in candidate_gap: + gap_sl.seg_label = prev_label + + +# ── Public API ──────────────────────────────────────────────────────────────── + +class SegmentationLabelDeriver: + """Derives one segmentation label per LayoutLine from token-level JATS annotations.""" + + def __init__(self, config: Optional[SegmentationConfig] = None) -> None: + self.config = config or SegmentationConfig() + + def derive_labels( + self, + layout_document: LayoutDocument, + annotated: JatsAnnotatedLayoutDocument, + ) -> Dict[int, str]: + """Return a mapping from line_id → segmentation label string. + + Uses `LayoutLineMeta.line_id` as the key so callers can look up labels + without holding LayoutLine references. + """ + seg_lines = [ + _SegLine(layout_line=line, line_index=idx) + for idx, line in enumerate(layout_document.iter_all_lines()) + ] + + # ── Tier 1: majority-vote from JATS token labels ── + for sl in seg_lines: + label = _majority_vote_label(sl.layout_line.tokens, annotated) + if label: + sl.seg_label = label + + # ── Tier 2: coordinate-based margin detection ── + page_meta_by_number = _get_page_meta_by_page_number(layout_document) + _tag_by_coordinates(seg_lines, page_meta_by_number, self.config) + + # ── Tier 3: heuristic passes ── + _clear_front_beyond_threshold( + seg_lines, self.config.front_max_start_line_index + ) + _find_missing_page_numbers(seg_lines) + _tag_headnotes_by_text_repetition( + seg_lines, self.config.page_header_max_first_line_index + ) + _merge_gap_lines( + seg_lines, + enabled_labels={SEG_FRONT, SEG_ANNEX, SEG_REFERENCES}, + enabled_tail_labels={SEG_ANNEX}, + ) + + # ── Default remaining untagged lines → body ── + for sl in seg_lines: + if sl.seg_label is None: + sl.seg_label = SEG_BODY + + # Build id(LayoutLine) → label mapping — safe because layout_document holds strong refs + result: Dict[int, str] = {} + for sl in seg_lines: + if sl.seg_label: + result[id(sl.layout_line)] = sl.seg_label + return result diff --git a/sciencebeam_parser/training/jats/text_normalizer.py b/sciencebeam_parser/training/jats/text_normalizer.py new file mode 100644 index 00000000..a71eb8e1 --- /dev/null +++ b/sciencebeam_parser/training/jats/text_normalizer.py @@ -0,0 +1,40 @@ +import unicodedata + +_LIGATURE_MAP = str.maketrans({ + 'ff': 'ff', + 'fi': 'fi', + 'fl': 'fl', + 'ffi': 'ffi', + 'ffl': 'ffl', + 'æ': 'ae', + 'œ': 'oe', +}) + +_DASH_MAP = str.maketrans({ + ch: '-' + for ch in '‐‑‒–—―' +}) + +_QUOTE_MAP = str.maketrans({ + '‘': "'", + '’': "'", + '“': '"', + '”': '"', +}) + +_SOFT_HYPHEN = '­' + + +def normalize_text(text: str) -> str: + """Light normalisation: ligatures, dashes, quotes, soft hyphens, NFC.""" + text = text.translate(_LIGATURE_MAP) + text = text.translate(_DASH_MAP) + text = text.translate(_QUOTE_MAP) + text = text.replace(_SOFT_HYPHEN, '') + return unicodedata.normalize('NFC', text) + + +def normalize_for_alignment(text: str) -> str: + """Aggressive normalisation for match scoring: lowercase + collapsed whitespace.""" + text = normalize_text(text) + return ' '.join(text.lower().split()) diff --git a/tests/training/jats/__init__.py b/tests/training/jats/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/training/jats/test_aligner.py b/tests/training/jats/test_aligner.py new file mode 100644 index 00000000..25bbeb4a --- /dev/null +++ b/tests/training/jats/test_aligner.py @@ -0,0 +1,369 @@ +from typing import Optional + +from sciencebeam_parser.document.layout_document import ( + LayoutBlock, + LayoutDocument, + LayoutLine, + LayoutPage, +) +from sciencebeam_parser.training.jats.aligner import AlignmentConfig, LayoutDocumentJatsAligner +from sciencebeam_parser.training.jats.field_extractor import JatsFieldValue +from sciencebeam_parser.training.jats.field_vocab import JatsFieldNames, JatsSubFieldNames + + +def _make_doc(*line_texts: str) -> LayoutDocument: + lines = [LayoutLine.for_text(t) for t in line_texts] + return LayoutDocument(pages=[LayoutPage(blocks=[LayoutBlock(lines=lines)])]) + + +def _fv(text: str, field: str = JatsFieldNames.BODY_SECTION_TITLE, sub: Optional[str] = None): + return JatsFieldValue(text=text, field_name=field, sub_field_name=sub) + + +class TestLayoutDocumentJatsAligner: + def _align(self, doc, field_values, **kwargs): + config = AlignmentConfig(**kwargs) if kwargs else None + return LayoutDocumentJatsAligner(config).align(doc, field_values) + + def test_empty_field_values_returns_unannotated(self): + doc = _make_doc('Hello world') + annotated = self._align(doc, []) + assert annotated.coverage_ratio() == 0.0 + + def test_exact_match_labels_all_tokens(self): + doc = _make_doc('Introduction') + annotated = self._align(doc, [_fv('Introduction')]) + tokens = list(doc.iter_all_tokens()) + assert all( + annotated.get_token_field(t) == JatsFieldNames.BODY_SECTION_TITLE + for t in tokens + ) + + def test_multi_token_match(self): + doc = _make_doc('The role of autophagy') + annotated = self._align(doc, [_fv('The role of autophagy')]) + tokens = list(doc.iter_all_tokens()) + assert all( + annotated.get_token_field(t) == JatsFieldNames.BODY_SECTION_TITLE + for t in tokens + ) + + def test_unmatched_tokens_have_no_label(self): + doc = _make_doc('Introduction', 'Some unrelated text here') + annotated = self._align(doc, [_fv('Introduction')]) + unrelated_tokens = list(doc.iter_all_lines())[1].tokens + assert all(annotated.get_token_field(t) is None for t in unrelated_tokens) + + def test_sub_field_overrides_parent_field(self): + doc = _make_doc('Smith J 2020') + fvs = [ + _fv('Smith J 2020', JatsFieldNames.REFERENCE), + _fv('Smith J', JatsFieldNames.REFERENCE, JatsSubFieldNames.REFERENCE_AUTHOR), + ] + annotated = self._align(doc, fvs) + tokens = list(doc.iter_all_tokens()) + # 'Smith' and 'J' tokens should have sub_field set + smith_token = next(t for t in tokens if t.text == 'Smith') + assert annotated.get_token_field(smith_token) == JatsFieldNames.REFERENCE + assert annotated.get_token_sub_field(smith_token) == JatsSubFieldNames.REFERENCE_AUTHOR + + def test_fuzzy_match_with_minor_variant(self): + # em-dash variant in haystack + doc = _make_doc('foo—bar baz') + annotated = self._align(doc, [_fv('foo-bar baz')]) + tokens = list(doc.iter_all_tokens()) + assert any(annotated.get_token_field(t) is not None for t in tokens) + + def test_coverage_ratio_partial(self): + doc = _make_doc('Title text', 'Other text') + annotated = self._align(doc, [_fv('Title text', JatsFieldNames.TITLE)]) + ratio = annotated.coverage_ratio() + assert 0 < ratio < 1.0 + + def test_coverage_ratio_full(self): + doc = _make_doc('only this') + annotated = self._align(doc, [_fv('only this')]) + assert annotated.coverage_ratio() == 1.0 + + def test_multiple_fields_labeled_correctly(self): + doc = _make_doc('My Title', 'John Smith') + fvs = [ + _fv('My Title', JatsFieldNames.TITLE), + _fv('John Smith', JatsFieldNames.AUTHOR), + ] + annotated = self._align(doc, fvs) + title_tokens = list(doc.iter_all_lines())[0].tokens + author_tokens = list(doc.iter_all_lines())[1].tokens + assert all(annotated.get_token_field(t) == JatsFieldNames.TITLE for t in title_tokens) + assert all(annotated.get_token_field(t) == JatsFieldNames.AUTHOR for t in author_tokens) + + def test_front_matter_author_found_before_abstract_in_haystack(self): + # Simulate the common PDF layout: author appears BEFORE the abstract in the + # document, but JATS XML lists the abstract first. With a long abstract the + # old last_match_end-based search_start would skip past the author's position; + # the front-matter region constraint fixes this. + abstract_text = ' '.join(['word'] * 60) # long enough that lme-200 > author pos + doc = _make_doc('Author Name', abstract_text, 'Introduction') + fvs = [ + _fv(abstract_text, JatsFieldNames.ABSTRACT), + _fv('Author Name', JatsFieldNames.AUTHOR), + _fv('Introduction', JatsFieldNames.BODY_SECTION_TITLE), + ] + annotated = self._align(doc, fvs) + author_tokens = list(doc.iter_all_lines())[0].tokens + intro_tokens = list(doc.iter_all_lines())[2].tokens + assert all(annotated.get_token_field(t) == JatsFieldNames.AUTHOR for t in author_tokens) + assert all( + annotated.get_token_field(t) == JatsFieldNames.BODY_SECTION_TITLE + for t in intro_tokens + ) + + def test_affiliation_found_after_body_when_not_in_front_matter(self): + # Some journals place full author/affiliation blocks at the END of the paper. + # The front-matter region constraint is soft: if the aff text is not found in + # [0, abstract_end + buffer] it should fall back to a global search. + abstract_text = ' '.join(['word'] * 60) + aff_text = 'Department of Computer Science University of Toronto Canada' + doc = _make_doc(abstract_text, 'Introduction Body', aff_text) + fvs = [ + _fv(abstract_text, JatsFieldNames.ABSTRACT), + _fv(aff_text, JatsFieldNames.AUTHOR_AFF), + ] + annotated = self._align(doc, fvs) + aff_tokens = list(doc.iter_all_lines())[2].tokens + assert all(annotated.get_token_field(t) == JatsFieldNames.AUTHOR_AFF for t in aff_tokens) + + def test_keywords_anchored_without_keywords_title(self): + # Even when there is no KEYWORDS_TITLE in the JATS (kwd-group has no <title>), + # the combined keyword string should not match in the abstract. The fallback + # anchors the search to just after front_matter_end (abstract end). + abstract_text = 'gene regulation and enhancers in limb development' + doc = _make_doc(abstract_text, 'gene regulation, enhancers') + fvs = [ + _fv(abstract_text, JatsFieldNames.ABSTRACT), + # no KEYWORDS_TITLE — single combined value per GROBID guidelines + _fv('gene regulation, enhancers', JatsFieldNames.KEYWORDS), + ] + annotated = self._align(doc, fvs) + kw_tokens = list(doc.iter_all_lines())[1].tokens + assert all(annotated.get_token_field(t) == JatsFieldNames.KEYWORDS for t in kw_tokens) + + def test_keywords_anchored_to_keywords_section(self): + # Per GROBID guidelines, all keywords form ONE field value. The whole list + # should be tagged on the keywords line, not in the abstract where the same + # words appear individually. + abstract_text = 'Bayesian confidence readout in dynamic stimuli tasks' + doc = _make_doc(abstract_text, 'Keywords confidence, Bayesian, DDM') + fvs = [ + _fv(abstract_text, JatsFieldNames.ABSTRACT), + _fv('Keywords', JatsFieldNames.KEYWORDS_TITLE), + _fv('confidence, Bayesian, DDM', JatsFieldNames.KEYWORDS), + ] + annotated = self._align(doc, fvs) + kw_tokens = list(doc.iter_all_lines())[1].tokens + kw_by_text = {t.text.lower().strip(','): t for t in kw_tokens} + assert annotated.get_token_field(kw_by_text['keywords']) == JatsFieldNames.KEYWORDS_TITLE + assert annotated.get_token_field(kw_by_text['confidence']) == JatsFieldNames.KEYWORDS + assert annotated.get_token_field(kw_by_text['bayesian']) == JatsFieldNames.KEYWORDS + assert annotated.get_token_field(kw_by_text['ddm']) == JatsFieldNames.KEYWORDS + + def test_sub_field_confined_to_parent_range(self): + # "Canada" appears in both lines; sub-field containment should pin the + # AUTHOR_AFF_COUNTRY match to the aff line, not the earlier occurrence. + aff_text = 'University of Toronto Canada' + # Put the ambiguous "Canada" earlier in the document than the aff block. + doc = _make_doc( + 'Some funding from Canada', + aff_text, + ) + fvs = [ + JatsFieldValue(aff_text, JatsFieldNames.AUTHOR_AFF), + JatsFieldValue( + 'Canada', JatsFieldNames.AUTHOR_AFF, + sub_field_name=JatsSubFieldNames.AUTHOR_AFF_COUNTRY, + ), + ] + annotated = self._align(doc, fvs) + aff_tokens = list(doc.iter_all_lines())[1].tokens + canada_in_aff = next(t for t in aff_tokens if 'canada' in t.text.lower()) + assert annotated.get_token_sub_field(canada_in_aff) == JatsSubFieldNames.AUTHOR_AFF_COUNTRY + + def test_consecutive_affiliations_have_distinct_instance_ids(self): + # Each JATS <aff> must produce a separate TEI <affiliation> element. The + # mechanism relies on the aligner assigning a distinct instance_id to each + # main (sub_field_name=None) AUTHOR_AFF field value so the header label fn + # emits B- on the first token of every new affiliation — even when no + # <address> tokens appear between them to force a label change. + aff1_text = '1 Institut Barcelona Spain' + aff2_text = '2 Cochrane Iberoamerica Madrid Spain' + doc = _make_doc(aff1_text, aff2_text) + fvs = [ + JatsFieldValue(aff1_text, JatsFieldNames.AUTHOR_AFF), + JatsFieldValue(aff2_text, JatsFieldNames.AUTHOR_AFF), + ] + annotated = self._align(doc, fvs) + aff1_tokens = list(doc.iter_all_lines())[0].tokens + aff2_tokens = list(doc.iter_all_lines())[1].tokens + assert all( + annotated.get_token_field(t) == JatsFieldNames.AUTHOR_AFF + for t in aff1_tokens + aff2_tokens + ) + # First aff → instance 1, second aff → instance 2: must differ + assert annotated.get_token_instance(aff1_tokens[0]) == 1 + assert annotated.get_token_instance(aff2_tokens[0]) == 2 + assert ( + annotated.get_token_instance(aff1_tokens[0]) + != annotated.get_token_instance(aff2_tokens[0]) + ) + + def test_abstract_does_not_label_sidebar_content(self): + # PDFs sometimes have a sidebar (e.g. Open Peer Review box) between the + # first and second page of an abstract. The JATS abstract needle contains + # no sidebar text, so Smith-Waterman creates only tiny (size ≤ 4) scatter + # blocks while traversing the sidebar. Those blocks must NOT cause sidebar + # tokens to be labelled as abstract. + page1 = ( + 'Every day important healthcare decisions are made with incomplete ' + 'information about the effects of the healthcare interventions ' + 'available. It is necessary to invest in strategies that allow access ' + 'to reliable and updated evidence on which to base health decisions.' + ) + # Sidebar: distinct proper-noun vocabulary with no 5+-char substring in page1/page2 + sidebar = ( + 'Open Peer Approval Ingrid Schmitt Lozano Maastricht ' + 'Pontificia Universidad Catolica panel assessment' + ) + page2 = ( + 'The project will be developed in three complementary phases. ' + 'Expected results include an effective capacity-building strategy ' + 'for health system organizations to implement the living evidence model.' + ) + abstract_text = page1 + ' ' + page2 + # Doc reading order: page1, then sidebar, then page2 (three separate lines) + doc = _make_doc(page1, sidebar, page2) + fvs = [_fv(abstract_text, JatsFieldNames.ABSTRACT)] + annotated = self._align(doc, fvs) + + lines = list(doc.iter_all_lines()) + sidebar_tokens = lines[1].tokens + labeled_sidebar = [ + t.text for t in sidebar_tokens + if annotated.get_token_field(t) == JatsFieldNames.ABSTRACT + ] + assert labeled_sidebar == [], ( + f'Sidebar tokens incorrectly labelled as abstract: {labeled_sidebar}' + ) + + # Page-2 abstract content (not the very first boundary token) must be labelled + page2_tokens = lines[2].tokens + page2_by_text = {t.text.lower(): t for t in page2_tokens} + for word in ('expected', 'capacity', 'building', 'organizations'): + if word in page2_by_text: + assert annotated.get_token_field(page2_by_text[word]) == JatsFieldNames.ABSTRACT, ( + f"Page-2 abstract token '{word}' was not labelled" + ) + + def test_doi_sub_field_all_tokens_labeled(self): + # The DOI tokenises into many short sub-tokens (2, 1, 4, 1, 3, 8 chars each). + # Without special handling the anchor+chain filter only labels the last long + # segment; all preceding dot/slash/digit segments should also be labeled. + ref_text = 'Smith J 2020 Some paper J Virol 10.1128/mBio.00524-13' + doi = '10.1128/mBio.00524-13' + doc = _make_doc(ref_text) + fvs = [ + _fv(ref_text, JatsFieldNames.REFERENCE), + _fv(doi, JatsFieldNames.REFERENCE, JatsSubFieldNames.REFERENCE_DOI), + ] + annotated = self._align(doc, fvs) + tokens = list(doc.iter_all_tokens()) + doi_tokens = [ + t for t in tokens + if annotated.get_token_sub_field(t) == JatsSubFieldNames.REFERENCE_DOI + ] + doi_text = ''.join(t.text for t in doi_tokens) + assert doi_text == doi, ( + f'Expected DOI tokens "{doi}", got "{doi_text}"' + ) + + def test_doi_sub_field_labeled_when_split_across_line_break(self): + # DOI split at end-of-line hyphen: PDF tokenizer emits a bare '-' as the last + # token of the first line, which the aligner strips (skip_tokens). All + # prefix sub-tokens before the join must still receive the DOI sub-field label. + ref_text = 'Smith J 2020 Some paper JRS 10.3233/JRS-201017' + doi = '10.3233/JRS-201017' + # Two-line doc: first line ends with the hyphen, second line has the suffix. + doc = _make_doc('Smith J 2020 Some paper JRS 10.3233/JRS-', '201017') + fvs = [ + _fv(ref_text, JatsFieldNames.REFERENCE), + _fv(doi, JatsFieldNames.REFERENCE, JatsSubFieldNames.REFERENCE_DOI), + ] + annotated = self._align(doc, fvs) + tokens = list(doc.iter_all_tokens()) + doi_tokens = [ + t for t in tokens + if annotated.get_token_sub_field(t) == JatsSubFieldNames.REFERENCE_DOI + ] + doi_labeled_text = ''.join(t.text for t in doi_tokens) + # tokens_in_range re-includes the bare '-' skip token because it follows the + # labeled 'JRS' token, so the full hyphenated form is reconstructed. + expected_labeled = '10.3233/JRS-201017' + assert doi_labeled_text == expected_labeled, ( + f'Expected labeled DOI text "{expected_labeled}", got "{doi_labeled_text}"' + ) + + def test_reference_spanning_page_break_does_not_label_headnote(self): + # A reference whose text spans a PDF page break has a running page header + # ("Journal Name 2025, 5:251") interleaved between its pre-break and + # post-break tokens. The SW blocks for a reference sub-field (e.g. article + # title) that crosses the break must NOT cause the headnote line to be + # labeled as a reference field. + # The anchor+chain filter handles this: the large gap between the last + # pre-break anchor block and the headnote blocks means they are not + # within_gap and are not part of pre_anchor, so they are dropped. + headnote = 'Journal Name 2025, 5:251 Last updated: 13 MAR 2026' + ref_text = ( + 'Smith J 2020 Clogging phenomenon in continuous casting of steel ' + 'a review Steel Res Int 10.1002/srin.201800' + ) + # Doc layout: ref pre-break line, then the running headnote, then ref suffix + doc = _make_doc( + 'Smith J 2020 Clogging phenomenon in continuous casting of steel', + headnote, + 'a review Steel Res Int 10.1002/srin.201800', + ) + fvs = [ + _fv(ref_text, JatsFieldNames.REFERENCE), + _fv( + 'Clogging phenomenon in continuous casting of steel a review', + JatsFieldNames.REFERENCE, + JatsSubFieldNames.REFERENCE_ARTICLE_TITLE, + ), + ] + annotated = self._align(doc, fvs) + lines = list(doc.iter_all_lines()) + headnote_tokens = lines[1].tokens + labeled_headnote = [ + t.text for t in headnote_tokens + if annotated.get_token_sub_field(t) == JatsSubFieldNames.REFERENCE_ARTICLE_TITLE + ] + assert labeled_headnote == [], ( + f'Headnote tokens incorrectly labeled as reference sub-field: {labeled_headnote}' + ) + + def test_heading_label_not_overwritten_by_paragraph_mid_token_match(self): + # Regression: SW alignment for a paragraph starting with "O crescimento" + # finds the last 'o' of "Introdução" (the heading token's tail char) as a + # spurious 2-char block ('o '). Without the token-boundary guard in the + # pre-anchor pass, tokens_in_range on that block returns the "Introdução" + # heading token and overwrites its BODY_SECTION_TITLE label with + # BODY_SECTION_PARAGRAPH. + doc = _make_doc('Introdução', 'O crescimento da pandemia do Covid') + fvs = [ + _fv('Introdução', JatsFieldNames.BODY_SECTION_TITLE), + _fv('O crescimento da pandemia do Covid', JatsFieldNames.BODY_SECTION_PARAGRAPH), + ] + annotated = self._align(doc, fvs) + heading_token = list(doc.iter_all_lines())[0].tokens[0] + assert annotated.get_token_field(heading_token) == JatsFieldNames.BODY_SECTION_TITLE, ( + f'Heading token label was overwritten to {annotated.get_token_field(heading_token)!r}' + ) diff --git a/tests/training/jats/test_coverage.py b/tests/training/jats/test_coverage.py new file mode 100644 index 00000000..98b5df3b --- /dev/null +++ b/tests/training/jats/test_coverage.py @@ -0,0 +1,111 @@ +from sciencebeam_parser.document.layout_document import ( + LayoutBlock, + LayoutDocument, + LayoutLine, + LayoutPage, + LayoutToken, +) +from sciencebeam_parser.training.jats.annotated_document import JatsAnnotatedLayoutDocument +from sciencebeam_parser.training.jats.coverage import CoverageResult, check_coverage +from sciencebeam_parser.training.jats.field_vocab import JatsFieldNames + + +def _make_annotated(*field_names: str) -> JatsAnnotatedLayoutDocument: + tokens = [LayoutToken(text=f'token_{i}') for i in range(len(field_names))] + line = LayoutLine(tokens=tokens) + doc = LayoutDocument(pages=[LayoutPage(blocks=[LayoutBlock(lines=[line])])]) + annotated = JatsAnnotatedLayoutDocument(layout_document=doc) + for token, field in zip(tokens, field_names): + if field: + annotated.set_token_label(token, field) + return annotated + + +class TestCoverageResult: + def test_passing_when_no_requirements(self): + result = CoverageResult() + assert result.is_passing + + def test_failing_when_required_field_missing(self): + result = CoverageResult(required_fields_missing={JatsFieldNames.TITLE}) + assert not result.is_passing + + def test_failing_when_matching_field_not_aligned(self): + result = CoverageResult( + required_matching_fields_missing={JatsFieldNames.ABSTRACT} + ) + assert not result.is_passing + + def test_str_ok(self): + assert str(CoverageResult()) == 'OK' + + def test_str_shows_missing_field(self): + result = CoverageResult(required_fields_missing={'title'}) + assert 'title' in str(result) + + +class TestCheckCoverage: + def test_required_field_present_and_aligned_passes(self): + annotated = _make_annotated(JatsFieldNames.TITLE) + result = check_coverage( + annotated=annotated, + field_values_by_field={JatsFieldNames.TITLE: True}, + required_fields=[JatsFieldNames.TITLE], + require_matching_fields=[], + ) + assert result.is_passing + assert JatsFieldNames.TITLE in result.required_fields_present + + def test_required_field_absent_from_jats_fails(self): + annotated = _make_annotated() + result = check_coverage( + annotated=annotated, + field_values_by_field={}, + required_fields=[JatsFieldNames.TITLE], + require_matching_fields=[], + ) + assert not result.is_passing + assert JatsFieldNames.TITLE in result.required_fields_missing + + def test_required_field_present_but_not_aligned_fails(self): + annotated = _make_annotated() # no labels assigned + result = check_coverage( + annotated=annotated, + field_values_by_field={JatsFieldNames.TITLE: True}, + required_fields=[JatsFieldNames.TITLE], + require_matching_fields=[], + ) + assert not result.is_passing + assert JatsFieldNames.TITLE in result.required_fields_missing + + def test_require_matching_field_absent_in_jats_passes(self): + annotated = _make_annotated() + result = check_coverage( + annotated=annotated, + field_values_by_field={}, # field not present in JATS + required_fields=[], + require_matching_fields=[JatsFieldNames.ABSTRACT], + ) + assert result.is_passing + + def test_require_matching_field_present_but_not_aligned_fails(self): + annotated = _make_annotated() # no labels + result = check_coverage( + annotated=annotated, + field_values_by_field={JatsFieldNames.ABSTRACT: True}, + required_fields=[], + require_matching_fields=[JatsFieldNames.ABSTRACT], + ) + assert not result.is_passing + assert JatsFieldNames.ABSTRACT in result.required_matching_fields_missing + + def test_require_matching_field_present_and_aligned_passes(self): + annotated = _make_annotated(JatsFieldNames.ABSTRACT) + result = check_coverage( + annotated=annotated, + field_values_by_field={JatsFieldNames.ABSTRACT: True}, + required_fields=[], + require_matching_fields=[JatsFieldNames.ABSTRACT], + ) + assert result.is_passing + assert JatsFieldNames.ABSTRACT in result.required_matching_fields_matched diff --git a/tests/training/jats/test_field_extractor.py b/tests/training/jats/test_field_extractor.py new file mode 100644 index 00000000..df48821d --- /dev/null +++ b/tests/training/jats/test_field_extractor.py @@ -0,0 +1,434 @@ +from collections import defaultdict + +from lxml import etree + +from sciencebeam_parser.training.jats.field_extractor import JatsFieldExtractor +from sciencebeam_parser.training.jats.field_vocab import JatsFieldNames, JatsSubFieldNames + + +def _parse_jats(xml: str) -> etree._Element: + return etree.fromstring(xml.encode()) + + +def _field_values_for(xml: str): + root = _parse_jats(xml) + return list(JatsFieldExtractor().iter_field_values(root)) + + +def _fields_by_name(values): + d = defaultdict(list) + for v in values: + d[v.field_name].append(v) + return d + + +class TestTitle: + def test_extracts_title(self): + fvs = _field_values_for( + '<article>' + '<front><article-meta><title-group>' + '<article-title>My Title</article-title>' + '</title-group></article-meta></front>' + '</article>' + ) + titles = [v for v in fvs if v.field_name == JatsFieldNames.TITLE] + assert len(titles) == 1 + assert titles[0].text == 'My Title' + assert titles[0].sub_field_name is None + + def test_no_title_gives_no_values(self): + fvs = _field_values_for('<article><front><article-meta></article-meta></front></article>') + titles = [v for v in fvs if v.field_name == JatsFieldNames.TITLE] + assert titles == [] + + +class TestAbstract: + def test_extracts_abstract(self): + fvs = _field_values_for( + '<article><front><article-meta>' + '<abstract><p>This is the abstract.</p></abstract>' + '</article-meta></front></article>' + ) + abstracts = [v for v in fvs if v.field_name == JatsFieldNames.ABSTRACT] + assert len(abstracts) == 1 + assert 'abstract' in abstracts[0].text.lower() + + +class TestAuthor: + def test_extracts_author_name(self): + fvs = _field_values_for( + '<article><front><article-meta>' + '<contrib-group>' + '<contrib contrib-type="author"><name>' + '<surname>Smith</surname><given-names>John</given-names>' + '</name></contrib>' + '</contrib-group>' + '</article-meta></front></article>' + ) + authors = [v for v in fvs if v.field_name == JatsFieldNames.AUTHOR] + assert len(authors) == 1 + assert authors[0].text == 'John Smith' + + def test_authors_merged_per_contrib_group(self): + # All authors in one <contrib-group> are emitted as a single AUTHOR field + # value so the aligner labels the full byline (including separators). + fvs = _field_values_for( + '<article><front><article-meta>' + '<contrib-group>' + '<contrib contrib-type="author">' + '<name><surname>Smith</surname><given-names>John</given-names></name>' + '<xref ref-type="aff">1</xref>' + '</contrib>' + '<contrib contrib-type="author">' + '<name><surname>Jones</surname><given-names>Mary</given-names></name>' + '<xref ref-type="aff">1</xref>' + '<xref ref-type="corresp">*</xref>' + '</contrib>' + '</contrib-group>' + '</article-meta></front></article>' + ) + authors = [v for v in fvs if v.field_name == JatsFieldNames.AUTHOR] + assert len(authors) == 1 + assert authors[0].text == 'John Smith 1 Mary Jones 1 *' + + def test_author_multiple_contrib_groups_emit_separately(self): + fvs = _field_values_for( + '<article><front><article-meta>' + '<contrib-group>' + '<contrib contrib-type="author">' + '<name><surname>Smith</surname><given-names>John</given-names></name>' + '</contrib>' + '</contrib-group>' + '<contrib-group>' + '<contrib contrib-type="author">' + '<name><surname>Jones</surname><given-names>Mary</given-names></name>' + '</contrib>' + '</contrib-group>' + '</article-meta></front></article>' + ) + authors = [v for v in fvs if v.field_name == JatsFieldNames.AUTHOR] + assert len(authors) == 2 + assert authors[0].text == 'John Smith' + assert authors[1].text == 'Mary Jones' + + +class TestKeywords: + def test_extracts_keyword_values(self): + fvs = _field_values_for( + '<article><front><article-meta>' + '<kwd-group>' + '<title>Keywords' + 'machine learning' + 'deep learning' + '' + '
' + ) + keywords = [v for v in fvs if v.field_name == JatsFieldNames.KEYWORDS] + assert len(keywords) == 1 + assert keywords[0].text == 'machine learning, deep learning' + + def test_keywords_title_extracted_for_segmentation(self): + # KEYWORDS_TITLE is emitted so the segmentation model can label the heading + # line as
. It is intentionally absent from HEADER_LABEL_BY_FIELD + # so the header model leaves the "Keywords" token unlabelled. + fvs = _field_values_for( + '
' + '' + 'Keywords' + 'machine learning' + '' + '
' + ) + kw_titles = [v for v in fvs if v.field_name == JatsFieldNames.KEYWORDS_TITLE] + assert len(kw_titles) == 1 + assert kw_titles[0].text == 'Keywords' + + +class TestAffiliation: + def test_extracts_affiliation_text(self): + fvs = _field_values_for( + '
' + 'MIT, Cambridge' + '
' + ) + affs = [v for v in fvs if v.field_name == JatsFieldNames.AUTHOR_AFF] + assert len(affs) >= 1 + assert 'MIT' in affs[0].text + + def test_extracts_aff_institution_subfield(self): + fvs = _field_values_for( + '
' + 'MIT' + '
' + ) + sub = [v for v in fvs if v.sub_field_name == JatsSubFieldNames.AUTHOR_AFF_INSTITUTION] + assert len(sub) == 1 + assert sub[0].text == 'MIT' + + def test_extracts_aff_country_subfield(self): + fvs = _field_values_for( + '
' + 'MITUSA' + '
' + ) + sub = [v for v in fvs if v.sub_field_name == JatsSubFieldNames.AUTHOR_AFF_COUNTRY] + assert len(sub) == 1 + assert sub[0].text == 'USA' + + def test_addr_bulk_includes_institution_tail_when_no_addr_line(self): + # Bioarxiv-style affs: UCL, GOSH, London WC1N 1EH, + # United Kingdom — city/postcode in institution tail, no . + # The AUTHOR_AFF_ADDR bulk value must cover "London WC1N 1EH United Kingdom" so + # those tokens get the
label rather than staying in . + fvs = _field_values_for( + '
' + '' + '' + 'UCL, GOSH' + ', London WC1N 1EH, ' + 'United Kingdom' + '' + '
' + ) + addr = [v for v in fvs if v.sub_field_name == JatsSubFieldNames.AUTHOR_AFF_ADDR] + assert len(addr) == 1 + assert 'London WC1N 1EH' in addr[0].text + assert 'United Kingdom' in addr[0].text + + def test_addr_bulk_empty_when_aff_fully_unstructured(self): + # Label-only affs (no , , ) cannot be split; + # no AUTHOR_AFF_ADDR value should be emitted. + fvs = _field_values_for( + '
' + 'Institut Barcelona Spain' + '
' + ) + addr = [v for v in fvs if v.sub_field_name == JatsSubFieldNames.AUTHOR_AFF_ADDR] + assert addr == [] + + def test_institution_tail_excluded_when_no_country_or_addr_line(self): + # Guard: institution tail is only address content when a or + # confirms the aff has structured address content. Without that anchor the tail + # may be continuation of the institution/department name: some publishers split a + # single department name across two tags, leaving the second half + # as the tail of the first — e.g. + # Dept of Microbiology, Immunology and Parasitology, + # University X, City, Country + # "Immunology and Parasitology" is NOT an address. + fvs = _field_values_for( + '
' + '' + '' + 'Departamento de Microbiologia' + ', Imunologia e Parasitologia, ' + 'Universidade Federal de Santa Catarina' + ', Florianopolis, Brasil' + '' + '
' + ) + addr = [v for v in fvs if v.sub_field_name == JatsSubFieldNames.AUTHOR_AFF_ADDR] + assert addr == [] + + +class TestReference: + def test_extracts_reference_text(self): + fvs = _field_values_for( + '
' + '' + 'A Study' + '2020' + '' + '
' + ) + refs = [v for v in fvs if v.field_name == JatsFieldNames.REFERENCE] + assert len(refs) >= 1 + assert 'A Study' in refs[0].text + + def test_extracts_reference_article_title_subfield(self): + fvs = _field_values_for( + '
' + '' + 'A Study' + '' + '
' + ) + sub = [v for v in fvs if v.sub_field_name == JatsSubFieldNames.REFERENCE_ARTICLE_TITLE] + assert len(sub) == 1 + assert sub[0].text == 'A Study' + + def test_extracts_reference_year_subfield(self): + fvs = _field_values_for( + '
' + '2020' + '
' + ) + sub = [v for v in fvs if v.sub_field_name == JatsSubFieldNames.REFERENCE_YEAR] + assert sub[0].text == '2020' + + def test_extracts_reference_doi_subfield(self): + fvs = _field_values_for( + '
' + '' + '10.1234/test' + '' + '
' + ) + sub = [v for v in fvs if v.sub_field_name == JatsSubFieldNames.REFERENCE_DOI] + assert len(sub) == 1 + assert sub[0].text == '10.1234/test' + + def test_extracts_reference_web_subfield_for_url_ext_link(self): + fvs = _field_values_for( + '
' + '' + '' + 'http://www.doi.org/10.5281/zenodo.6647010' + '' + '' + '
' + ) + sub = [v for v in fvs if v.sub_field_name == JatsSubFieldNames.REFERENCE_WEB] + assert len(sub) == 1 + assert sub[0].text == 'http://www.doi.org/10.5281/zenodo.6647010' + + def test_reference_web_ignores_reference_source_ext_link(self): + fvs = _field_values_for( + '
' + '' + 'Reference Source' + '' + '
' + ) + sub = [v for v in fvs if v.sub_field_name == JatsSubFieldNames.REFERENCE_WEB] + assert len(sub) == 0 + + +class TestBodySections: + def test_extracts_body_section_title(self): + fvs = _field_values_for( + '
Introduction' + '

Some text.

' + ) + titles = [v for v in fvs if v.field_name == JatsFieldNames.BODY_SECTION_TITLE] + assert len(titles) == 1 + assert titles[0].text == 'Introduction' + + def test_extracts_body_paragraph(self): + fvs = _field_values_for( + '
Intro' + '

Paragraph text.

' + ) + paras = [v for v in fvs if v.field_name == JatsFieldNames.BODY_SECTION_PARAGRAPH] + assert len(paras) == 1 + assert 'Paragraph' in paras[0].text + + +class TestAcknowledgement: + def test_extracts_ack_paragraph(self): + fvs = _field_values_for( + '

We thank everyone.

' + ) + ack = [v for v in fvs if v.field_name == JatsFieldNames.ACK_SECTION_PARAGRAPH] + assert len(ack) == 1 + assert 'thank' in ack[0].text + + +class TestFunding: + def test_extracts_funding_statement(self): + fvs = _field_values_for( + '
' + '' + 'Supported by grant 123.' + '' + '
' + ) + funding = [v for v in fvs if v.field_name == JatsFieldNames.FUNDING] + assert len(funding) == 1 + assert 'grant 123' in funding[0].text + + def test_extracts_multiple_funding_statements(self): + fvs = _field_values_for( + '
' + '' + 'Grant A funded this.' + 'The funder had no role.' + '' + '
' + ) + funding = [v for v in fvs if v.field_name == JatsFieldNames.FUNDING] + assert len(funding) == 2 + + +class TestCopyright: + def test_extracts_copyright_statement(self): + fvs = _field_values_for( + '
' + '' + 'Copyright 2022 Author et al.' + '' + '
' + ) + cr = [v for v in fvs if v.field_name == JatsFieldNames.COPYRIGHT] + assert len(cr) == 1 + assert '2022' in cr[0].text + + def test_extracts_license_paragraph(self): + fvs = _field_values_for( + '
' + '' + '' + 'Open access under CC-BY 4.0.' + '' + '' + '
' + ) + cr = [v for v in fvs if v.field_name == JatsFieldNames.COPYRIGHT] + assert len(cr) == 1 + assert 'CC-BY' in cr[0].text + + +class TestSubArticle: + def test_extracts_sub_article_paragraphs(self): + fvs = _field_values_for( + '
' + '' + '

This manuscript is well written.

' + '' + '
' + ) + sub = [v for v in fvs if v.field_name == JatsFieldNames.SUB_ARTICLE] + assert len(sub) == 1 + assert 'well written' in sub[0].text + + def test_extracts_sub_article_titles(self): + fvs = _field_values_for( + '
' + '' + 'Reviewer Report' + '

Some comments.

' + '' + '
' + ) + sub = [v for v in fvs if v.field_name == JatsFieldNames.SUB_ARTICLE] + texts = [v.text for v in sub] + assert any('Reviewer Report' in t for t in texts) + assert any('comments' in t for t in texts) + + def test_main_article_body_not_labeled_as_sub_article(self): + fvs = _field_values_for( + '
' + 'Introduction' + '

Main article text.

' + '' + '

Review text.

' + '' + '
' + ) + sub = [v for v in fvs if v.field_name == JatsFieldNames.SUB_ARTICLE] + body = [v for v in fvs if v.field_name == JatsFieldNames.BODY_SECTION_PARAGRAPH] + assert len(sub) == 1 + assert 'Review text' in sub[0].text + assert len(body) == 1 + assert 'Main article' in body[0].text diff --git a/tests/training/jats/test_segmentation.py b/tests/training/jats/test_segmentation.py new file mode 100644 index 00000000..e0a8e380 --- /dev/null +++ b/tests/training/jats/test_segmentation.py @@ -0,0 +1,209 @@ +from sciencebeam_parser.document.layout_document import ( + LayoutBlock, + LayoutDocument, + LayoutLine, + LayoutPage, + LayoutPageCoordinates, + LayoutPageMeta, + LayoutToken, +) +from sciencebeam_parser.training.jats.annotated_document import JatsAnnotatedLayoutDocument +from sciencebeam_parser.training.jats.field_vocab import JatsFieldNames +from sciencebeam_parser.training.jats.segmentation import ( + SEG_BODY, + SEG_FRONT, + SEG_HEADNOTE, + SEG_PAGE, + SEG_REFERENCES, + SegmentationConfig, + SegmentationLabelDeriver, +) + + +def _make_page_meta(page_number: int = 1, height: float = 1000.0) -> LayoutPageMeta: + return LayoutPageMeta( + page_number=page_number, + coordinates=LayoutPageCoordinates( + x=0, y=0, width=600, height=height, page_number=page_number + ), + ) + + +def _make_token( + text: str, + page_number: int = 1, + y: float = 500.0, +) -> LayoutToken: + return LayoutToken( + text=text, + coordinates=LayoutPageCoordinates( + x=10, y=y, width=50, height=12, page_number=page_number + ), + ) + + +def _make_line(*texts: str, y: float = 500.0, page_number: int = 1) -> LayoutLine: + tokens = [_make_token(t, page_number=page_number, y=y) for t in texts] + return LayoutLine(tokens=tokens) + + +def _make_doc_with_page( + *blocks: LayoutBlock, page_height: float = 1000.0, page_number: int = 1 +) -> LayoutDocument: + page_meta = _make_page_meta(page_number=page_number, height=page_height) + page = LayoutPage(blocks=list(blocks), meta=page_meta) + return LayoutDocument(pages=[page]) + + +def _annotate(doc: LayoutDocument, field_by_line_index: dict) -> JatsAnnotatedLayoutDocument: + annotated = JatsAnnotatedLayoutDocument(layout_document=doc) + lines = list(doc.iter_all_lines()) + for line_idx, field_name in field_by_line_index.items(): + for token in lines[line_idx].tokens: + annotated.set_token_label(token, field_name) + return annotated + + +def _derive_labels(doc, annotated, **config_kwargs): + config = SegmentationConfig(**config_kwargs) if config_kwargs else None + return SegmentationLabelDeriver(config).derive_labels(doc, annotated) + + +class TestMajorityVoteLabeling: + def test_title_tokens_give_header_label(self): + line = _make_line('My', 'Title') + doc = _make_doc_with_page(LayoutBlock(lines=[line])) + annotated = _annotate(doc, {0: JatsFieldNames.TITLE}) + labels = _derive_labels(doc, annotated) + assert labels[id(line)] == SEG_FRONT + + def test_body_paragraph_gives_body_label(self): + line = _make_line('Some', 'body', 'text') + doc = _make_doc_with_page(LayoutBlock(lines=[line])) + annotated = _annotate(doc, {0: JatsFieldNames.BODY_SECTION_PARAGRAPH}) + labels = _derive_labels(doc, annotated) + assert labels[id(line)] == SEG_BODY + + def test_reference_tokens_give_references_label(self): + line = _make_line('Smith', '2020') + doc = _make_doc_with_page(LayoutBlock(lines=[line])) + annotated = _annotate(doc, {0: JatsFieldNames.REFERENCE}) + labels = _derive_labels(doc, annotated) + assert labels[id(line)] == SEG_REFERENCES + + def test_unannotated_line_defaults_to_body(self): + line = _make_line('Unknown', 'content') + doc = _make_doc_with_page(LayoutBlock(lines=[line])) + annotated = JatsAnnotatedLayoutDocument(layout_document=doc) + labels = _derive_labels(doc, annotated) + assert labels[id(line)] == SEG_BODY + + def test_majority_vote_mixed_line(self): + # 3 tokens labeled as TITLE, 1 as BODY_SECTION_PARAGRAPH → should give FRONT (header) + line = _make_line('My', 'Title', 'Here', 'text') + doc = _make_doc_with_page(LayoutBlock(lines=[line])) + annotated = JatsAnnotatedLayoutDocument(layout_document=doc) + tokens = line.tokens + for t in tokens[:3]: + annotated.set_token_label(t, JatsFieldNames.TITLE) + annotated.set_token_label(tokens[3], JatsFieldNames.BODY_SECTION_PARAGRAPH) + labels = _derive_labels(doc, annotated) + assert labels[id(line)] == SEG_FRONT + + +class TestCoordinateBasedDetection: + def test_line_at_top_of_page_becomes_headnote(self): + line = _make_line('Running', 'header', y=20.0) # 20/1000 = 2% < 8% + doc = _make_doc_with_page(LayoutBlock(lines=[line]), page_height=1000.0) + annotated = JatsAnnotatedLayoutDocument(layout_document=doc) + labels = _derive_labels(doc, annotated, headnote_y_ratio=0.08) + assert labels[id(line)] == SEG_HEADNOTE + + def test_line_in_middle_of_page_is_not_headnote(self): + line = _make_line('Normal', 'content', y=500.0) # 50% of page + doc = _make_doc_with_page(LayoutBlock(lines=[line]), page_height=1000.0) + annotated = JatsAnnotatedLayoutDocument(layout_document=doc) + labels = _derive_labels(doc, annotated) + assert labels[id(line)] != SEG_HEADNOTE + + def test_numeric_line_at_bottom_becomes_page(self): + line = _make_line('42', y=950.0) # 95% of page > 92% + doc = _make_doc_with_page(LayoutBlock(lines=[line]), page_height=1000.0) + annotated = JatsAnnotatedLayoutDocument(layout_document=doc) + labels = _derive_labels(doc, annotated, footnote_y_ratio=0.92) + assert labels[id(line)] == SEG_PAGE + + +class TestGapMerge: + def test_untagged_line_between_front_lines_gets_front(self): + line_front1 = _make_line('Title', 'text', y=200.0) + line_gap = _make_line('Some', 'untagged', 'stuff', y=220.0) + line_front2 = _make_line('More', 'header', y=240.0) + block = LayoutBlock(lines=[line_front1, line_gap, line_front2]) + doc = _make_doc_with_page(block) + annotated = JatsAnnotatedLayoutDocument(layout_document=doc) + for t in line_front1.tokens: + annotated.set_token_label(t, JatsFieldNames.TITLE) + for t in line_front2.tokens: + annotated.set_token_label(t, JatsFieldNames.AUTHOR) + labels = _derive_labels(doc, annotated) + assert labels[id(line_front1)] == SEG_FRONT + assert labels[id(line_front2)] == SEG_FRONT + assert labels[id(line_gap)] == SEG_FRONT + + def test_untagged_line_after_body_stays_body(self): + line_body = _make_line('Body', 'paragraph', y=400.0) + line_gap = _make_line('More', 'stuff', y=420.0) + block = LayoutBlock(lines=[line_body, line_gap]) + doc = _make_doc_with_page(block) + annotated = JatsAnnotatedLayoutDocument(layout_document=doc) + for t in line_body.tokens: + annotated.set_token_label(t, JatsFieldNames.BODY_SECTION_PARAGRAPH) + labels = _derive_labels(doc, annotated) + assert labels[id(line_body)] == SEG_BODY + # gap after body → body (default) + assert labels[id(line_gap)] == SEG_BODY + + +class TestFrontThreshold: + def test_front_block_starting_within_threshold_is_kept(self): + # A front block starting at line 0 should not be cleared + lines = [_make_line(f'word{i}') for i in range(5)] + block = LayoutBlock(lines=lines) + doc = _make_doc_with_page(block) + annotated = JatsAnnotatedLayoutDocument(layout_document=doc) + for t in lines[0].tokens: + annotated.set_token_label(t, JatsFieldNames.TITLE) + for t in lines[4].tokens: + annotated.set_token_label(t, JatsFieldNames.AUTHOR) + labels = _derive_labels(doc, annotated, front_max_start_line_index=80) + assert labels[id(lines[0])] == SEG_FRONT + assert labels[id(lines[4])] == SEG_FRONT + + def test_front_block_starting_beyond_threshold_is_cleared(self): + # A front block starting at line 100 (> 80) should be cleared → defaults to body + lines = [_make_line(f'word{i}') for i in range(3)] + block = LayoutBlock(lines=lines) + doc = _make_doc_with_page(block) + annotated = JatsAnnotatedLayoutDocument(layout_document=doc) + # Annotate only the last line as front — its block starts at index 2 which is + # not cleared. Use a config with a very low threshold to test the clearing logic. + for t in lines[2].tokens: + annotated.set_token_label(t, JatsFieldNames.AUTHOR_NOTES) + labels = _derive_labels(doc, annotated, front_max_start_line_index=1) + # line 2 starts its own front block at index 2 > threshold 1 → cleared → body + assert labels[id(lines[2])] == SEG_BODY + + +class TestTextRepetitionHeadnote: + def test_repeated_line_near_top_becomes_headnote(self): + # No coordinates → will fall through to text-repetition detection + lines_no_coords = [LayoutLine(tokens=[LayoutToken(text=t)]) for t in ['Journal Name'] * 3] + block2 = LayoutBlock(lines=lines_no_coords) + doc = LayoutDocument(pages=[LayoutPage(blocks=[block2])]) + annotated = JatsAnnotatedLayoutDocument(layout_document=doc) + labels = _derive_labels( + doc, annotated, page_header_max_first_line_index=10 + ) + for line in lines_no_coords: + assert labels.get(id(line)) == SEG_HEADNOTE diff --git a/tests/training/jats/test_text_normalizer.py b/tests/training/jats/test_text_normalizer.py new file mode 100644 index 00000000..318c724f --- /dev/null +++ b/tests/training/jats/test_text_normalizer.py @@ -0,0 +1,51 @@ +from sciencebeam_parser.training.jats.text_normalizer import ( + normalize_text, + normalize_for_alignment, +) + + +class TestNormalizeText: + def test_passthrough_ascii(self): + assert normalize_text('hello world') == 'hello world' + + def test_ligature_fi(self): + assert normalize_text('figure') == 'figure' + + def test_ligature_fl(self): + assert normalize_text('flow') == 'flow' + + def test_ligature_ff(self): + assert normalize_text('off') == 'off' + + def test_em_dash_to_hyphen(self): + assert normalize_text('foo—bar') == 'foo-bar' + + def test_en_dash_to_hyphen(self): + assert normalize_text('foo–bar') == 'foo-bar' + + def test_curly_quotes(self): + assert normalize_text('‘it’s') == "'it's" + + def test_double_curly_quotes(self): + assert normalize_text('“hello”') == '"hello"' + + def test_soft_hyphen_removed(self): + # U+00AD soft hyphen + assert normalize_text('hyp­hen') == 'hyphen' + + +class TestNormalizeForAlignment: + def test_lowercase(self): + assert normalize_for_alignment('Hello World') == 'hello world' + + def test_collapses_whitespace(self): + assert normalize_for_alignment('foo \t bar') == 'foo bar' + + def test_strips_leading_trailing(self): + assert normalize_for_alignment(' hello ') == 'hello' + + def test_applies_ligature_normalisation(self): + assert normalize_for_alignment('figure') == 'figure' + + def test_applies_dash_normalisation(self): + assert normalize_for_alignment('foo—bar') == 'foo-bar'