Generate JATS-guided training data without model cascade dependency

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 \

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)

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()

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 <output-path>/<split>/<corpus>/.  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()

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("<xml/>")
+            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()