From 73e904756ec544e1dfd3a0ee0847f2df474d50b7 Mon Sep 17 00:00:00 2001
From: Daniel Ecer <de-code@users.noreply.github.com>
Date: Fri, 26 Jun 2026 13:56:02 +0100
Subject: [PATCH 01/12] Expand JATS reference splitting to whole-line level

Two improvements to JATS-guided training data generation for the
reference-segmenter and citation models:

1. Line-level gap fill in _split_references_by_jats_instance: unlabeled
   lines (e.g. DOI continuation lines that the aligner didn't match)
   inherit the preceding line's REFERENCE instance_id rather than being
   silently dropped. Restores DOI continuation text that was absent from
   citation training data.

2. Whole-line expansion in ReferenceSegmenterModelTrainingDataGenerator
   .get_jats_label_fn: if any token on a PDF line is labeled REFERENCE,
   all tokens on that line receive the same <reference>/<label> tag.
   Eliminates the fragmented <bibl> spans caused by unlabeled separator
   tokens (colons, semicolons, whitespace between sub-fields).

Both changes rest on the invariant that no two distinct references share
a single PDF line, which holds for all observed training documents.
---
 .../training/cli/generate_data.py             | 135 ++++++-
 tests/training/cli/generate_data_test.py      | 333 +++++++++++++++++-
 2 files changed, 460 insertions(+), 8 deletions(-)

diff --git a/sciencebeam_parser/training/cli/generate_data.py b/sciencebeam_parser/training/cli/generate_data.py
index 01cbfebb..8ffe71db 100644
--- a/sciencebeam_parser/training/cli/generate_data.py
+++ b/sciencebeam_parser/training/cli/generate_data.py
@@ -17,7 +17,7 @@
 
 from sciencebeam_parser.utils.io import glob, makedirs, write_bytes, write_text
 
-from sciencebeam_parser.document.layout_document import LayoutDocument
+from sciencebeam_parser.document.layout_document import LayoutBlock, LayoutDocument, LayoutLine
 from sciencebeam_parser.document.semantic_document import (
     SemanticMixedContentWrapper,
     SemanticRawAffiliationAddress,
@@ -52,6 +52,7 @@
     CITATION_LABEL_BY_SUB_FIELD,
     FULLTEXT_LABEL_BY_FIELD,
     HEADER_LABEL_BY_FIELD,
+    JatsFieldNames,
     JatsSubFieldNames,
 )
 from sciencebeam_parser.training.jats.field_extractor import JatsFieldExtractor
@@ -766,8 +767,6 @@ def iter_model_layout_documents(
         layout_document: LayoutDocument,
         document_context: TrainingDataDocumentContext
     ) -> Iterable[LayoutDocument]:
-        reference_segmenter_model = document_context.fulltext_models.reference_segmenter_model
-        citation_model = document_context.fulltext_models.citation_model
         segmentation_label_result = get_segmentation_label_result(
             layout_document,
             document_context=document_context
@@ -775,6 +774,42 @@ def iter_model_layout_documents(
         references_layout_document = segmentation_label_result.get_filtered_document_by_label(
             '<references>'
         ).remove_empty_blocks()
+        annotated = document_context.jats_annotated_document
+        if annotated is not None:
+            # Group author tokens by reference instance_id, preserving line structure.
+            # Each instance_id becomes one LayoutBlock (one per reference).
+            author_tokens_by_instance: Dict[int, List] = {}
+            for token in references_layout_document.iter_all_tokens():
+                entry = annotated.token_label_by_id.get(id(token))
+                if (
+                    entry
+                    and entry[0] == JatsFieldNames.REFERENCE
+                    and entry[1] == JatsSubFieldNames.REFERENCE_AUTHOR
+                ):
+                    inst = entry[2]
+                    if inst not in author_tokens_by_instance:
+                        author_tokens_by_instance[inst] = []
+                    author_tokens_by_instance[inst].append(token)
+            if not author_tokens_by_instance:
+                return []
+            LOGGER.info(
+                'found author tokens for %d references (JATS path)',
+                len(author_tokens_by_instance)
+            )
+            return [LayoutDocument.for_blocks([
+                LayoutBlock.for_tokens(tokens)
+                for tokens in author_tokens_by_instance.values()
+            ])]
+        # Fallback: reference segmenter + citation model
+        return self._iter_via_models(references_layout_document, document_context)
+
+    def _iter_via_models(
+        self,
+        references_layout_document: LayoutDocument,
+        document_context: TrainingDataDocumentContext,
+    ) -> Iterable[LayoutDocument]:
+        reference_segmenter_model = document_context.fulltext_models.reference_segmenter_model
+        citation_model = document_context.fulltext_models.citation_model
         reference_segmenter_labeled_layout_tokens = (
             get_labeled_layout_tokens_for_model_and_layout_document(
                 model=reference_segmenter_model,
@@ -819,7 +854,6 @@ def iter_model_layout_documents(
         LOGGER.info('semantic_raw_author_list count: %d', len(semantic_raw_author_list))
         if not semantic_raw_author_list:
             return []
-
         return [
             LayoutDocument.for_blocks([
                 block
@@ -945,10 +979,97 @@ def iter_model_layout_documents(
         ]
 
 
+def _split_references_by_jats_instance(
+    refs_layout_doc: LayoutDocument,
+    annotated: JatsAnnotatedLayoutDocument,
+) -> List[LayoutDocument]:
+    """Split a references sub-document into one LayoutDocument per JATS <ref> instance.
+
+    Pass 1: assign each line to a REFERENCE instance_id (majority vote over labeled tokens).
+    Lines with no labeled tokens inherit the preceding line's instance (gap fill) — this
+    covers DOI continuation lines and other unlabeled continuation text.
+
+    Pass 2: assign each block to the dominant instance_id among its lines.
+    Blocks where no line has any instance (e.g. the list title) are skipped.
+    """
+    # Pass 1: line-level instance assignment with gap fill
+    line_instance: Dict[int, int] = {}
+    last_instance: Optional[int] = None
+    for line in refs_layout_doc.iter_all_lines():
+        instance_counts: Dict[int, int] = {}
+        for token in line.tokens:
+            entry = annotated.token_label_by_id.get(id(token))
+            if entry and entry[0] == JatsFieldNames.REFERENCE:
+                inst = entry[2]
+                instance_counts[inst] = instance_counts.get(inst, 0) + 1
+        if instance_counts:
+            dominant = max(instance_counts, key=instance_counts.__getitem__)
+            line_instance[id(line)] = dominant
+            last_instance = dominant
+        elif last_instance is not None:
+            line_instance[id(line)] = last_instance
+
+    # Pass 2: block-level assignment via line votes
+    blocks_by_instance: Dict[int, List[LayoutBlock]] = {}
+    for block in refs_layout_doc.iter_all_blocks():
+        instance_votes: Dict[int, int] = {}
+        for line in block.lines:
+            line_inst = line_instance.get(id(line))
+            if line_inst is not None:
+                instance_votes[line_inst] = instance_votes.get(line_inst, 0) + 1
+        if not instance_votes:
+            continue
+        dominant = max(instance_votes, key=instance_votes.__getitem__)
+        if dominant not in blocks_by_instance:
+            blocks_by_instance[dominant] = []
+        blocks_by_instance[dominant].append(block)
+    return [
+        LayoutDocument.for_blocks(blocks)
+        for blocks in blocks_by_instance.values()
+    ]
+
+
 class ReferenceSegmenterModelTrainingDataGenerator(AbstractDocumentModelTrainingDataGenerator):
     def get_main_model(self, document_context: TrainingDataDocumentContext) -> Model:
         return document_context.fulltext_models.reference_segmenter_model
 
+    def get_jats_label_fn(self) -> Optional[JatsLabelFn]:
+        line_label_cache: Dict[int, Optional[str]] = {}
+
+        def _line_label(
+            line: LayoutLine, annotated: JatsAnnotatedLayoutDocument
+        ) -> Optional[str]:
+            has_label_token = False
+            has_reference_token = False
+            for token in line.tokens:
+                entry = annotated.token_label_by_id.get(id(token))
+                if not entry or entry[0] != JatsFieldNames.REFERENCE:
+                    continue
+                if entry[1] == JatsSubFieldNames.REFERENCE_LABEL:
+                    has_label_token = True
+                else:
+                    has_reference_token = True
+            if has_label_token:
+                return '<label>'
+            if has_reference_token:
+                return '<reference>'
+            return None
+
+        def fn(
+            annotated: JatsAnnotatedLayoutDocument,
+            _seg_labels: Dict[int, str],
+            md: LayoutModelData,
+        ) -> Optional[str]:
+            line = md.layout_line
+            if line is None or not md.layout_token:
+                return None
+            line_id = id(line)
+            if line_id not in line_label_cache:
+                line_label_cache[line_id] = _line_label(line, annotated)
+            return line_label_cache[line_id]
+
+        return fn
+
     def iter_model_layout_documents(
         self,
         layout_document: LayoutDocument,
@@ -988,7 +1109,6 @@ def iter_model_layout_documents(
         layout_document: LayoutDocument,
         document_context: TrainingDataDocumentContext
     ) -> Iterable[LayoutDocument]:
-        reference_segmenter_model = document_context.fulltext_models.reference_segmenter_model
         segmentation_label_result = get_segmentation_label_result(
             layout_document,
             document_context=document_context
@@ -996,6 +1116,11 @@ def iter_model_layout_documents(
         references_layout_document = segmentation_label_result.get_filtered_document_by_label(
             '<references>'
         ).remove_empty_blocks()
+        annotated = document_context.jats_annotated_document
+        if annotated is not None:
+            return _split_references_by_jats_instance(references_layout_document, annotated)
+        # Fallback: reference segmenter model
+        reference_segmenter_model = document_context.fulltext_models.reference_segmenter_model
         reference_segmenter_labeled_layout_tokens = (
             get_labeled_layout_tokens_for_model_and_layout_document(
                 model=reference_segmenter_model,
diff --git a/tests/training/cli/generate_data_test.py b/tests/training/cli/generate_data_test.py
index 5fe00b26..976295ae 100644
--- a/tests/training/cli/generate_data_test.py
+++ b/tests/training/cli/generate_data_test.py
@@ -3,7 +3,7 @@
 import re
 import gzip
 from pathlib import Path
-from typing import Iterable, Iterator, Optional, Sequence
+from typing import Dict, Iterable, Iterator, List, Optional, Sequence
 from unittest.mock import MagicMock, patch
 
 import pytest
@@ -15,8 +15,11 @@
 from sciencebeam_parser.document.layout_document import (
     LayoutBlock,
     LayoutDocument,
-    LayoutPage
+    LayoutLine,
+    LayoutPage,
+    join_layout_tokens,
 )
+from sciencebeam_parser.models.data import LayoutModelData
 from sciencebeam_parser.document.tei.common import get_tei_xpath_text_content_list
 from sciencebeam_parser.models.data import DEFAULT_DOCUMENT_FEATURES_CONTEXT
 from sciencebeam_parser.models.training_data import TeiTrainingDataGenerator
@@ -43,10 +46,18 @@
 from sciencebeam_parser.models.citation.training_data import (
     CitationTeiTrainingDataGenerator
 )
+from sciencebeam_parser.training.jats.annotated_document import JatsAnnotatedLayoutDocument
+from sciencebeam_parser.training.jats.field_vocab import JatsFieldNames, JatsSubFieldNames
 import sciencebeam_parser.training.cli.generate_data as generate_data_module
 from sciencebeam_parser.training.cli.generate_data import (
+    CitationModelTrainingDataGenerator,
+    ModelResultCache,
+    NameCitationModelTrainingDataGenerator,
+    ReferenceSegmenterModelTrainingDataGenerator,
+    TrainingDataDocumentContext,
+    _split_references_by_jats_instance,
     generate_training_data_for_layout_document,
-    main
+    main,
 )
 
 from tests.processors.fulltext.model_mocks import MockFullTextModels
@@ -655,3 +666,319 @@ def test_should_allow_to_use_directory_structure(
                 suffix
             )
             assert file_path.exists()
+
+
+# ── Helpers for JATS-based tests ─────────────────────────────────────────────
+
+
+def _doc_text(doc: LayoutDocument) -> str:
+    return join_layout_tokens(list(doc.iter_all_tokens()))
+
+
+def _annotate_block_as_reference(
+    block: LayoutBlock,
+    annotated: JatsAnnotatedLayoutDocument,
+    instance_id: int,
+    sub_field: Optional[str] = None,
+) -> None:
+    for token in block.iter_all_tokens():
+        annotated.set_token_label(
+            token,
+            JatsFieldNames.REFERENCE,
+            sub_field_name=sub_field,
+            instance_id=instance_id,
+        )
+
+
+def _make_jats_context(
+    annotated: JatsAnnotatedLayoutDocument,
+    ref_blocks: List[LayoutBlock],
+) -> TrainingDataDocumentContext:
+    seg_labels: Dict[int, str] = {
+        id(line): '<references>'
+        for block in ref_blocks
+        for line in block.lines
+    }
+    return TrainingDataDocumentContext(
+        output_path='/tmp/test',
+        source_filename='test.pdf',
+        document_features_context=DEFAULT_DOCUMENT_FEATURES_CONTEXT,
+        fulltext_models=MagicMock(),
+        use_model=False,
+        use_directory_structure=False,
+        model_result_cache=ModelResultCache(),
+        gzip_enabled=False,
+        jats_annotated_document=annotated,
+        jats_segmentation_labels=seg_labels,
+    )
+
+
+# ── _split_references_by_jats_instance ────────────────────────────────────────
+
+@log_on_exception
+class TestSplitReferencesByJatsInstance:
+    def test_returns_one_document_per_reference_instance(self):
+        ref1_block = LayoutBlock.for_text('Smith 2020 Title Journal')
+        ref2_block = LayoutBlock.for_text('Jones 2019 Article Nature')
+        refs_doc = LayoutDocument(pages=[LayoutPage(blocks=[ref1_block, ref2_block])])
+        annotated = JatsAnnotatedLayoutDocument(layout_document=refs_doc)
+        _annotate_block_as_reference(ref1_block, annotated, instance_id=1)
+        _annotate_block_as_reference(ref2_block, annotated, instance_id=2)
+
+        result = _split_references_by_jats_instance(refs_doc, annotated)
+
+        assert len(result) == 2
+        assert _doc_text(result[0]) == ref1_block.text
+        assert _doc_text(result[1]) == ref2_block.text
+
+    def test_skips_blocks_with_no_reference_tokens(self):
+        title_block = LayoutBlock.for_text('References')
+        ref1_block = LayoutBlock.for_text('Smith 2020')
+        refs_doc = LayoutDocument(pages=[LayoutPage(blocks=[title_block, ref1_block])])
+        annotated = JatsAnnotatedLayoutDocument(layout_document=refs_doc)
+        _annotate_block_as_reference(ref1_block, annotated, instance_id=1)
+
+        result = _split_references_by_jats_instance(refs_doc, annotated)
+
+        assert len(result) == 1
+        assert _doc_text(result[0]) == ref1_block.text
+
+    def test_returns_empty_list_when_no_reference_annotations(self):
+        block = LayoutBlock.for_text('Some text')
+        refs_doc = LayoutDocument(pages=[LayoutPage(blocks=[block])])
+        annotated = JatsAnnotatedLayoutDocument(layout_document=refs_doc)
+
+        result = _split_references_by_jats_instance(refs_doc, annotated)
+
+        assert not result
+
+    def test_assigns_mixed_block_to_dominant_instance(self):
+        # A block where most tokens belong to instance 1, one token to instance 2
+        block = LayoutBlock.for_text('Smith Jones Brown Extra')
+        refs_doc = LayoutDocument(pages=[LayoutPage(blocks=[block])])
+        annotated = JatsAnnotatedLayoutDocument(layout_document=refs_doc)
+        tokens = list(block.iter_all_tokens())
+        for t in tokens[:3]:
+            annotated.set_token_label(t, JatsFieldNames.REFERENCE, instance_id=1)
+        annotated.set_token_label(tokens[3], JatsFieldNames.REFERENCE, instance_id=2)
+
+        result = _split_references_by_jats_instance(refs_doc, annotated)
+
+        assert len(result) == 1  # dominant instance 1 wins, block not split
+
+    def test_gap_fill_includes_unlabeled_continuation_block(self):
+        # Simulates a DOI continuation line: ref1 block has labeled tokens, the
+        # continuation block (e.g. the second line of a wrapped DOI) has none.
+        # It should be attached to ref1, not silently dropped.
+        ref1_block = LayoutBlock.for_text('Smith 2020 doi')
+        cont_block = LayoutBlock.for_text('10.1234/continuation')
+        ref2_block = LayoutBlock.for_text('Jones 2019 Article Nature')
+        refs_doc = LayoutDocument(
+            pages=[LayoutPage(blocks=[ref1_block, cont_block, ref2_block])]
+        )
+        annotated = JatsAnnotatedLayoutDocument(layout_document=refs_doc)
+        _annotate_block_as_reference(ref1_block, annotated, instance_id=1)
+        _annotate_block_as_reference(ref2_block, annotated, instance_id=2)
+        # cont_block is intentionally left unannotated
+
+        result = _split_references_by_jats_instance(refs_doc, annotated)
+
+        assert len(result) == 2
+        doc1_text = _doc_text(result[0])
+        assert 'Smith' in doc1_text
+        assert '10.1234/continuation' in doc1_text  # gap-filled into ref1
+        assert 'Jones' not in doc1_text
+        doc2_text = _doc_text(result[1])
+        assert 'Jones' in doc2_text
+        assert '10.1234/continuation' not in doc2_text
+
+    def test_leading_unlabeled_block_is_still_skipped(self):
+        # A block appearing BEFORE any labeled content has no preceding instance
+        # to gap-fill from, so it should be dropped (e.g. the "References" title).
+        title_block = LayoutBlock.for_text('References')
+        ref1_block = LayoutBlock.for_text('Smith 2020')
+        refs_doc = LayoutDocument(
+            pages=[LayoutPage(blocks=[title_block, ref1_block])]
+        )
+        annotated = JatsAnnotatedLayoutDocument(layout_document=refs_doc)
+        _annotate_block_as_reference(ref1_block, annotated, instance_id=1)
+
+        result = _split_references_by_jats_instance(refs_doc, annotated)
+
+        assert len(result) == 1
+        assert 'References' not in _doc_text(result[0])
+        assert 'Smith' in _doc_text(result[0])
+
+
+# ── ReferenceSegmenterModelTrainingDataGenerator JATS label fn ───────────────
+
+
+def _make_md(line: LayoutLine, token_idx: int = 0) -> LayoutModelData:
+    return LayoutModelData(
+        data_line='token feats',
+        layout_line=line,
+        layout_token=line.tokens[token_idx],
+    )
+
+
+@log_on_exception
+class TestReferenceSegmenterJatsLabelFn:
+    def test_labels_whole_line_as_reference_when_any_token_labeled(self):
+        line = LayoutLine.for_text('Smith 2020 : Journal')
+        refs_doc = LayoutDocument(pages=[LayoutPage(blocks=[LayoutBlock(lines=[line])])])
+        annotated = JatsAnnotatedLayoutDocument(layout_document=refs_doc)
+        # Label only the first token
+        annotated.set_token_label(line.tokens[0], JatsFieldNames.REFERENCE, instance_id=1)
+
+        label_fn = ReferenceSegmenterModelTrainingDataGenerator().get_jats_label_fn()
+        assert label_fn is not None
+
+        # First token (labeled) → '<reference>'
+        assert label_fn(annotated, {}, _make_md(line, 0)) == '<reference>'
+        # Unlabeled token on the same line → also '<reference>'
+        assert label_fn(annotated, {}, _make_md(line, 2)) == '<reference>'
+
+    def test_returns_none_for_line_with_no_reference_tokens(self):
+        line = LayoutLine.for_text('Publisher Full Text')
+        refs_doc = LayoutDocument(pages=[LayoutPage(blocks=[LayoutBlock(lines=[line])])])
+        annotated = JatsAnnotatedLayoutDocument(layout_document=refs_doc)
+
+        label_fn = ReferenceSegmenterModelTrainingDataGenerator().get_jats_label_fn()
+        assert label_fn is not None
+
+        assert label_fn(annotated, {}, _make_md(line, 0)) is None
+
+    def test_labels_line_as_label_when_reference_label_token_present(self):
+        line = LayoutLine.for_text('[1] Smith 2020')
+        refs_doc = LayoutDocument(pages=[LayoutPage(blocks=[LayoutBlock(lines=[line])])])
+        annotated = JatsAnnotatedLayoutDocument(layout_document=refs_doc)
+        annotated.set_token_label(
+            line.tokens[0], JatsFieldNames.REFERENCE,
+            sub_field_name=JatsSubFieldNames.REFERENCE_LABEL, instance_id=1,
+        )
+        # remaining tokens are plain REFERENCE
+        for t in line.tokens[1:]:
+            annotated.set_token_label(t, JatsFieldNames.REFERENCE, instance_id=1)
+
+        label_fn = ReferenceSegmenterModelTrainingDataGenerator().get_jats_label_fn()
+        assert label_fn is not None
+
+        # All tokens on the line get '<label>' because at least one is REFERENCE_LABEL
+        for idx in range(len(line.tokens)):
+            assert label_fn(annotated, {}, _make_md(line, idx)) == '<label>'
+
+
+# ── CitationModelTrainingDataGenerator JATS path ──────────────────────────────
+
+@log_on_exception
+class TestCitationModelJatsPath:
+    def test_uses_jats_instance_split_and_skips_reference_segmenter(self):
+        ref1_block = LayoutBlock.for_text('Smith 2020 A Title')
+        ref2_block = LayoutBlock.for_text('Jones 2019 B Journal')
+        layout_document = LayoutDocument(pages=[LayoutPage(blocks=[ref1_block, ref2_block])])
+        annotated = JatsAnnotatedLayoutDocument(layout_document=layout_document)
+        _annotate_block_as_reference(ref1_block, annotated, instance_id=1)
+        _annotate_block_as_reference(ref2_block, annotated, instance_id=2)
+        context = _make_jats_context(annotated, [ref1_block, ref2_block])
+
+        result = list(
+            CitationModelTrainingDataGenerator().iter_model_layout_documents(
+                layout_document, context
+            )
+        )
+
+        assert len(result) == 2
+        assert _doc_text(result[0]) == ref1_block.text
+        assert _doc_text(result[1]) == ref2_block.text
+        context.fulltext_models.reference_segmenter_model.assert_not_called()
+
+    def test_returns_empty_when_no_references_annotated(self):
+        block = LayoutBlock.for_text('Some text')
+        layout_document = LayoutDocument(pages=[LayoutPage(blocks=[block])])
+        annotated = JatsAnnotatedLayoutDocument(layout_document=layout_document)
+        context = _make_jats_context(annotated, [block])
+
+        result = list(
+            CitationModelTrainingDataGenerator().iter_model_layout_documents(
+                layout_document, context
+            )
+        )
+
+        assert not result
+
+
+# ── NameCitationModelTrainingDataGenerator JATS path ─────────────────────────
+
+@log_on_exception
+class TestNameCitationModelJatsPath:
+    def test_returns_author_tokens_without_running_models(self):
+        author_block = LayoutBlock.for_text('Smith J Jones M')
+        other_block = LayoutBlock.for_text('Title of Paper Journal 2020')
+        layout_document = LayoutDocument(
+            pages=[LayoutPage(blocks=[author_block, other_block])]
+        )
+        annotated = JatsAnnotatedLayoutDocument(layout_document=layout_document)
+        _annotate_block_as_reference(
+            author_block, annotated, instance_id=1,
+            sub_field=JatsSubFieldNames.REFERENCE_AUTHOR
+        )
+        _annotate_block_as_reference(other_block, annotated, instance_id=1)
+        context = _make_jats_context(annotated, [author_block, other_block])
+
+        result = list(
+            NameCitationModelTrainingDataGenerator().iter_model_layout_documents(
+                layout_document, context
+            )
+        )
+
+        assert len(result) == 1
+        result_text = _doc_text(result[0])
+        assert 'Smith' in result_text
+        assert 'Jones' in result_text
+        # Non-author tokens should not be present
+        assert 'Title' not in result_text
+        context.fulltext_models.reference_segmenter_model.assert_not_called()
+        context.fulltext_models.citation_model.assert_not_called()
+
+    def test_returns_empty_when_no_author_annotations(self):
+        block = LayoutBlock.for_text('Smith 2020')
+        layout_document = LayoutDocument(pages=[LayoutPage(blocks=[block])])
+        annotated = JatsAnnotatedLayoutDocument(layout_document=layout_document)
+        _annotate_block_as_reference(block, annotated, instance_id=1)
+        context = _make_jats_context(annotated, [block])
+
+        result = list(
+            NameCitationModelTrainingDataGenerator().iter_model_layout_documents(
+                layout_document, context
+            )
+        )
+
+        assert not result
+
+    def test_collects_authors_from_multiple_references(self):
+        author1_block = LayoutBlock.for_text('Smith J')
+        author2_block = LayoutBlock.for_text('Jones M')
+        layout_document = LayoutDocument(
+            pages=[LayoutPage(blocks=[author1_block, author2_block])]
+        )
+        annotated = JatsAnnotatedLayoutDocument(layout_document=layout_document)
+        _annotate_block_as_reference(
+            author1_block, annotated, instance_id=1,
+            sub_field=JatsSubFieldNames.REFERENCE_AUTHOR
+        )
+        _annotate_block_as_reference(
+            author2_block, annotated, instance_id=2,
+            sub_field=JatsSubFieldNames.REFERENCE_AUTHOR
+        )
+        context = _make_jats_context(annotated, [author1_block, author2_block])
+
+        result = list(
+            NameCitationModelTrainingDataGenerator().iter_model_layout_documents(
+                layout_document, context
+            )
+        )
+
+        assert len(result) == 1
+        result_text = _doc_text(result[0])
+        assert 'Smith' in result_text
+        assert 'Jones' in result_text

From f843fc2bcca098e115b537d4df3de4aff0d2912a Mon Sep 17 00:00:00 2001
From: Daniel Ecer <de-code@users.noreply.github.com>
Date: Fri, 26 Jun 2026 17:40:13 +0100
Subject: [PATCH 02/12] Fix reference-segmenter training data bibl boundaries
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Each reference in the generated TEI now lands in its own <bibl>,
with its label (e.g. "9.") at the start rather than stranded at the
end of the previous reference. This means the reference-segmenter
model trains on correctly structured data, improving its ability to
detect where one reference ends and the next begins — particularly
for references whose number appears on the same line as their content.

Also fixes block splitting for the citation model path: a layout block
that straddles two JATS instances is now split at the boundary rather
than assigned wholesale to one reference, and running headers/page
numbers inside the references section no longer bleed into adjacent
references via gap fill.
---
 .../training/cli/generate_data.py             | 121 +++++++++-----
 tests/training/cli/generate_data_test.py      | 158 +++++++++++++++++-
 2 files changed, 234 insertions(+), 45 deletions(-)

diff --git a/sciencebeam_parser/training/cli/generate_data.py b/sciencebeam_parser/training/cli/generate_data.py
index 889adfeb..b8510816 100644
--- a/sciencebeam_parser/training/cli/generate_data.py
+++ b/sciencebeam_parser/training/cli/generate_data.py
@@ -1,6 +1,7 @@
 # pylint: disable=too-many-lines
 from abc import ABC, abstractmethod
 import argparse
+from itertools import groupby
 import logging
 import os
 import multiprocessing
@@ -1013,23 +1014,25 @@ def iter_model_layout_documents(
         ]
 
 
-def _split_references_by_jats_instance(
+_GAP_FILL_SKIP_SEG_LABELS = frozenset({'<headnote>', '<page>', '<footnote>'})
+
+
+def _assign_line_instances(
     refs_layout_doc: LayoutDocument,
     annotated: JatsAnnotatedLayoutDocument,
-) -> List[LayoutDocument]:
-    """Split a references sub-document into one LayoutDocument per JATS <ref> instance.
-
-    Pass 1: assign each line to a REFERENCE instance_id (majority vote over labeled tokens).
-    Lines with no labeled tokens inherit the preceding line's instance (gap fill) — this
-    covers DOI continuation lines and other unlabeled continuation text.
+    seg_labels: Dict[int, str],
+) -> Dict[int, int]:
+    """Return id(line) → instance_id for lines that belong to a JATS REFERENCE instance.
 
-    Pass 2: assign each block to the dominant instance_id among its lines.
-    Blocks where no line has any instance (e.g. the list title) are skipped.
+    Lines with no labeled tokens inherit the preceding line's instance (gap fill).
+    Lines whose seg label is in _GAP_FILL_SKIP_SEG_LABELS are excluded from both
+    assignment and gap-fill carry.
     """
-    # Pass 1: line-level instance assignment with gap fill
     line_instance: Dict[int, int] = {}
     last_instance: Optional[int] = None
     for line in refs_layout_doc.iter_all_lines():
+        if seg_labels.get(id(line)) in _GAP_FILL_SKIP_SEG_LABELS:
+            continue
         instance_counts: Dict[int, int] = {}
         for token in line.tokens:
             entry = annotated.token_label_by_id.get(id(token))
@@ -1042,21 +1045,34 @@ def _split_references_by_jats_instance(
             last_instance = dominant
         elif last_instance is not None:
             line_instance[id(line)] = last_instance
+    return line_instance
+
+
+def _split_references_by_jats_instance(
+    refs_layout_doc: LayoutDocument,
+    annotated: JatsAnnotatedLayoutDocument,
+    jats_seg_labels: Optional[Dict[int, str]] = None,
+) -> List[LayoutDocument]:
+    """Split a references sub-document into one LayoutDocument per JATS <ref> instance.
 
-    # Pass 2: block-level assignment via line votes
+    Uses line-level instance assignment (with gap fill for unlabeled continuation lines)
+    followed by groupby-based block splitting. Blocks that contain lines from two different
+    instances are split at the instance boundary. Lines with headnote/page/footnote seg
+    labels and lines before any labeled content are excluded.
+    """
+    line_instance = _assign_line_instances(
+        refs_layout_doc, annotated, jats_seg_labels or {}
+    )
     blocks_by_instance: Dict[int, List[LayoutBlock]] = {}
     for block in refs_layout_doc.iter_all_blocks():
-        instance_votes: Dict[int, int] = {}
-        for line in block.lines:
-            line_inst = line_instance.get(id(line))
-            if line_inst is not None:
-                instance_votes[line_inst] = instance_votes.get(line_inst, 0) + 1
-        if not instance_votes:
-            continue
-        dominant = max(instance_votes, key=instance_votes.__getitem__)
-        if dominant not in blocks_by_instance:
-            blocks_by_instance[dominant] = []
-        blocks_by_instance[dominant].append(block)
+        for instance_id, line_group in groupby(
+            block.lines, key=lambda l: line_instance.get(id(l))
+        ):
+            if instance_id is None:
+                continue
+            if instance_id not in blocks_by_instance:
+                blocks_by_instance[instance_id] = []
+            blocks_by_instance[instance_id].append(LayoutBlock(lines=list(line_group)))
     return [
         LayoutDocument.for_blocks(blocks)
         for blocks in blocks_by_instance.values()
@@ -1070,25 +1086,16 @@ def get_main_model(self, document_context: TrainingDataDocumentContext) -> Model
         return document_context.fulltext_models.reference_segmenter_model
 
     def get_jats_label_fn(self) -> Optional[JatsLabelFn]:
-        line_label_cache: Dict[int, Optional[str]] = {}
+        line_first_instance_cache: Dict[int, Optional[int]] = {}
+        last_instance_id: Optional[int] = None
 
-        def _line_label(
+        def _get_line_first_ref_instance_id(
             line: LayoutLine, annotated: JatsAnnotatedLayoutDocument
-        ) -> Optional[str]:
-            has_label_token = False
-            has_reference_token = False
+        ) -> Optional[int]:
             for token in line.tokens:
                 entry = annotated.token_label_by_id.get(id(token))
-                if not entry or entry[0] != JatsFieldNames.REFERENCE:
-                    continue
-                if entry[1] == JatsSubFieldNames.REFERENCE_LABEL:
-                    has_label_token = True
-                else:
-                    has_reference_token = True
-            if has_label_token:
-                return '<label>'
-            if has_reference_token:
-                return '<reference>'
+                if entry and entry[0] == JatsFieldNames.REFERENCE:
+                    return entry[2]
             return None
 
         def fn(
@@ -1096,13 +1103,39 @@ def fn(
             _seg_labels: Dict[int, str],
             md: LayoutModelData,
         ) -> Optional[str]:
+            nonlocal last_instance_id
+            token = md.layout_token
             line = md.layout_line
-            if line is None or not md.layout_token:
+            if line is None or token is None:
                 return None
+            entry = annotated.token_label_by_id.get(id(token))
+            if entry and entry[0] == JatsFieldNames.REFERENCE:
+                cur_instance_id = entry[2]
+                is_label = entry[1] == JatsSubFieldNames.REFERENCE_LABEL
+                if cur_instance_id != last_instance_id and last_instance_id is not None:
+                    # Instance transition: emit B-prefixed label so the TEI generator
+                    # creates a new <bibl>.  B-<label> also fires the reset mechanism
+                    # (backs up to <listBibl> before opening the new <bibl>), which is
+                    # how the <label> subelement lands in the correct <bibl>.
+                    last_instance_id = cur_instance_id
+                    return 'B-<label>' if is_label else 'B-<reference>'
+                last_instance_id = cur_instance_id
+                return '<label>' if is_label else '<reference>'
+            # Unlabeled token: expand to <reference> if the line has any reference tokens.
+            # If the line's first annotated reference belongs to a new instance, claim
+            # the whole line for that instance by emitting B-<reference>.
             line_id = id(line)
-            if line_id not in line_label_cache:
-                line_label_cache[line_id] = _line_label(line, annotated)
-            return line_label_cache[line_id]
+            if line_id not in line_first_instance_cache:
+                line_first_instance_cache[line_id] = _get_line_first_ref_instance_id(
+                    line, annotated
+                )
+            line_instance_id = line_first_instance_cache[line_id]
+            if line_instance_id is None:
+                return None
+            if line_instance_id != last_instance_id and last_instance_id is not None:
+                last_instance_id = line_instance_id
+                return 'B-<reference>'
+            return '<reference>'
 
         return fn
 
@@ -1156,7 +1189,11 @@ def iter_model_layout_documents(
         ).remove_empty_blocks()
         annotated = document_context.jats_annotated_document
         if annotated is not None:
-            return _split_references_by_jats_instance(references_layout_document, annotated)
+            return _split_references_by_jats_instance(
+                references_layout_document,
+                annotated,
+                jats_seg_labels=document_context.jats_segmentation_labels,
+            )
         # Fallback: reference segmenter model
         reference_segmenter_model = document_context.fulltext_models.reference_segmenter_model
         reference_segmenter_labeled_layout_tokens = (
diff --git a/tests/training/cli/generate_data_test.py b/tests/training/cli/generate_data_test.py
index 976295ae..d531d743 100644
--- a/tests/training/cli/generate_data_test.py
+++ b/tests/training/cli/generate_data_test.py
@@ -1,3 +1,4 @@
+# pylint: disable=too-many-lines
 import logging
 import os
 import re
@@ -809,6 +810,49 @@ def test_leading_unlabeled_block_is_still_skipped(self):
         assert 'References' not in _doc_text(result[0])
         assert 'Smith' in _doc_text(result[0])
 
+    def test_splits_block_at_instance_boundary(self):
+        # Two references share one LayoutBlock (Burguete/Carvalho pattern):
+        # line 1 belongs to instance 1, line 2 belongs to instance 2.
+        # The block must be split so each reference gets its own sub-document.
+        line1 = LayoutLine.for_text('Burguete 2020 Title Journal')
+        line2 = LayoutLine.for_text('Carvalho 2019 Another Paper')
+        shared_block = LayoutBlock(lines=[line1, line2])
+        refs_doc = LayoutDocument(pages=[LayoutPage(blocks=[shared_block])])
+        annotated = JatsAnnotatedLayoutDocument(layout_document=refs_doc)
+        for token in line1.tokens:
+            annotated.set_token_label(token, JatsFieldNames.REFERENCE, instance_id=1)
+        for token in line2.tokens:
+            annotated.set_token_label(token, JatsFieldNames.REFERENCE, instance_id=2)
+
+        result = _split_references_by_jats_instance(refs_doc, annotated)
+
+        assert len(result) == 2
+        assert 'Burguete' in _doc_text(result[0])
+        assert 'Carvalho' not in _doc_text(result[0])
+        assert 'Carvalho' in _doc_text(result[1])
+        assert 'Burguete' not in _doc_text(result[1])
+
+    def test_excludes_headnote_lines_from_gap_fill(self):
+        # A page header that slipped into the references sub-document must not
+        # be included in any reference's sub-document.
+        ref1_block = LayoutBlock.for_text('Smith 2020 doi')
+        header_line = LayoutLine.for_text('Journal Name | Volume 1 | 2020')
+        header_block = LayoutBlock(lines=[header_line])
+        ref2_block = LayoutBlock.for_text('Jones 2019 Article')
+        refs_doc = LayoutDocument(
+            pages=[LayoutPage(blocks=[ref1_block, header_block, ref2_block])]
+        )
+        annotated = JatsAnnotatedLayoutDocument(layout_document=refs_doc)
+        _annotate_block_as_reference(ref1_block, annotated, instance_id=1)
+        _annotate_block_as_reference(ref2_block, annotated, instance_id=2)
+        jats_seg_labels = {id(header_line): '<headnote>'}
+
+        result = _split_references_by_jats_instance(refs_doc, annotated, jats_seg_labels)
+
+        assert len(result) == 2
+        assert 'Journal Name' not in _doc_text(result[0])
+        assert 'Journal Name' not in _doc_text(result[1])
+
 
 # ── ReferenceSegmenterModelTrainingDataGenerator JATS label fn ───────────────
 
@@ -863,9 +907,117 @@ def test_labels_line_as_label_when_reference_label_token_present(self):
         label_fn = ReferenceSegmenterModelTrainingDataGenerator().get_jats_label_fn()
         assert label_fn is not None
 
-        # All tokens on the line get '<label>' because at least one is REFERENCE_LABEL
-        for idx in range(len(line.tokens)):
-            assert label_fn(annotated, {}, _make_md(line, idx)) == '<label>'
+        # Only the REFERENCE_LABEL token itself gets '<label>'; the rest get '<reference>'
+        assert label_fn(annotated, {}, _make_md(line, 0)) == '<label>'
+        for idx in range(1, len(line.tokens)):
+            assert label_fn(annotated, {}, _make_md(line, idx)) == '<reference>'
+
+    def test_unlabeled_token_on_reference_line_gets_reference_not_label(self):
+        # "10. Author Name..." where "10." is REFERENCE_LABEL and the rest are plain REFERENCE.
+        # An additional unlabeled token on the same line should expand to '<reference>',
+        # NOT '<label>'.
+        line = LayoutLine.for_text('10. Author Name unlabeled')
+        refs_doc = LayoutDocument(pages=[LayoutPage(blocks=[LayoutBlock(lines=[line])])])
+        annotated = JatsAnnotatedLayoutDocument(layout_document=refs_doc)
+        annotated.set_token_label(
+            line.tokens[0], JatsFieldNames.REFERENCE,
+            sub_field_name=JatsSubFieldNames.REFERENCE_LABEL, instance_id=1,
+        )
+        annotated.set_token_label(line.tokens[1], JatsFieldNames.REFERENCE, instance_id=1)
+        annotated.set_token_label(line.tokens[2], JatsFieldNames.REFERENCE, instance_id=1)
+        # tokens[3] ("unlabeled") has no annotation
+
+        label_fn = ReferenceSegmenterModelTrainingDataGenerator().get_jats_label_fn()
+        assert label_fn is not None
+
+        assert label_fn(annotated, {}, _make_md(line, 0)) == '<label>'      # "10."
+        assert label_fn(annotated, {}, _make_md(line, 1)) == '<reference>'  # "Author"
+        assert label_fn(annotated, {}, _make_md(line, 2)) == '<reference>'  # "Name"
+        assert label_fn(annotated, {}, _make_md(line, 3)) == '<reference>'  # "unlabeled" expanded
+
+    def test_plain_reference_transition_emits_b_prefix_for_bibl_boundary(self):
+        # When a plain <reference> token starts a new instance (no <label>),
+        # the label fn returns 'B-<reference>' so the TEI generator creates a new
+        # <bibl> without losing the token.
+        line1 = LayoutLine.for_text('Smith 2020')
+        line2 = LayoutLine.for_text('Doe 2019')
+        refs_doc = LayoutDocument(pages=[LayoutPage(blocks=[
+            LayoutBlock(lines=[line1, line2])
+        ])])
+        annotated = JatsAnnotatedLayoutDocument(layout_document=refs_doc)
+        for t in line1.tokens:
+            annotated.set_token_label(t, JatsFieldNames.REFERENCE, instance_id=1)
+        for t in line2.tokens:
+            annotated.set_token_label(t, JatsFieldNames.REFERENCE, instance_id=2)
+
+        label_fn = ReferenceSegmenterModelTrainingDataGenerator().get_jats_label_fn()
+        assert label_fn is not None
+
+        # All of line 1 → '<reference>' (first instance, no prior instance)
+        for idx in range(len(line1.tokens)):
+            assert label_fn(annotated, {}, _make_md(line1, idx)) == '<reference>'
+        # First token of line 2 → 'B-<reference>' (B-prefix creates new bibl)
+        assert label_fn(annotated, {}, _make_md(line2, 0)) == 'B-<reference>'
+        # Remaining tokens of line 2 → '<reference>' (same instance, no transition)
+        for idx in range(1, len(line2.tokens)):
+            assert label_fn(annotated, {}, _make_md(line2, idx)) == '<reference>'
+
+    def test_unlabeled_token_on_new_instance_line_emits_b_reference(self):
+        # "9. Moraes R." — "9." is unlabeled but "Moraes R." is annotated as instance 2.
+        # The label fn should emit 'B-<reference>' for "9." so it claims the whole line
+        # for the new bibl, rather than appending "9." to the previous bibl.
+        line1 = LayoutLine.for_text('https://some.url/paper')
+        line2 = LayoutLine.for_text('9. Moraes R. Title')
+        refs_doc = LayoutDocument(pages=[LayoutPage(blocks=[
+            LayoutBlock(lines=[line1, line2])
+        ])])
+        annotated = JatsAnnotatedLayoutDocument(layout_document=refs_doc)
+        for t in line1.tokens:
+            annotated.set_token_label(t, JatsFieldNames.REFERENCE, instance_id=1)
+        # "9." (line2.tokens[0]) is intentionally NOT annotated (unlabeled in JATS)
+        for t in line2.tokens[1:]:
+            annotated.set_token_label(t, JatsFieldNames.REFERENCE, instance_id=2)
+
+        label_fn = ReferenceSegmenterModelTrainingDataGenerator().get_jats_label_fn()
+        assert label_fn is not None
+
+        for idx in range(len(line1.tokens)):
+            label_fn(annotated, {}, _make_md(line1, idx))  # advance state
+        # "9." is unlabeled but its line's first annotated token is instance 2 → B-prefix
+        assert label_fn(annotated, {}, _make_md(line2, 0)) == 'B-<reference>'
+        # Remaining tokens (annotated as instance 2) → same instance, no transition
+        for idx in range(1, len(line2.tokens)):
+            assert label_fn(annotated, {}, _make_md(line2, idx)) == '<reference>'
+
+    def test_label_token_at_instance_transition_emits_b_label(self):
+        # When a REFERENCE_LABEL token starts a new instance the label fn returns
+        # 'B-<label>' so the reset mechanism (which fires only on B-prefix) creates a
+        # new <bibl> with the label text correctly placed inside it.
+        line1 = LayoutLine.for_text('Smith 2020')
+        line2 = LayoutLine.for_text('[2] Doe 2019')
+        refs_doc = LayoutDocument(pages=[LayoutPage(blocks=[
+            LayoutBlock(lines=[line1, line2])
+        ])])
+        annotated = JatsAnnotatedLayoutDocument(layout_document=refs_doc)
+        for t in line1.tokens:
+            annotated.set_token_label(t, JatsFieldNames.REFERENCE, instance_id=1)
+        annotated.set_token_label(
+            line2.tokens[0], JatsFieldNames.REFERENCE,
+            sub_field_name=JatsSubFieldNames.REFERENCE_LABEL, instance_id=2,
+        )
+        for t in line2.tokens[1:]:
+            annotated.set_token_label(t, JatsFieldNames.REFERENCE, instance_id=2)
+
+        label_fn = ReferenceSegmenterModelTrainingDataGenerator().get_jats_label_fn()
+        assert label_fn is not None
+
+        for idx in range(len(line1.tokens)):
+            label_fn(annotated, {}, _make_md(line1, idx))  # advance state
+        # REFERENCE_LABEL at instance transition → 'B-<label>' (triggers reset + new bibl)
+        assert label_fn(annotated, {}, _make_md(line2, 0)) == 'B-<label>'
+        # Remaining tokens of line 2 → '<reference>'
+        for idx in range(1, len(line2.tokens)):
+            assert label_fn(annotated, {}, _make_md(line2, idx)) == '<reference>'
 
 
 # ── CitationModelTrainingDataGenerator JATS path ──────────────────────────────

From 7be87fc2caef578527a653518ee58807c8b832ba Mon Sep 17 00:00:00 2001
From: Daniel Ecer <de-code@users.noreply.github.com>
Date: Fri, 26 Jun 2026 19:57:40 +0100
Subject: [PATCH 03/12] Fix reference label and DOI continuation in JATS
 training data
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

  Bracket-style reference labels like `[1]`, `[2]` were not appearing in `<label>` elements in the
  generated TEI training data. The PDF tokeniser splits `[1]` into three tokens `[`, `1`, `]`,
  producing `[ 1 ]` in the alignment haystack. The Smith-Waterman library's traceback logic
  terminates early at gap moves, so it matched only the final `]` (quality 0.33, below threshold).
  A new `_try_bracket_label_match` fallback strips the brackets, locates the inner number via exact
  token match, and extends the range to include the surrounding bracket tokens. It also searches a
  small buffer before the nominal segment start, since the parent SW match often begins at `]` —
  leaving `[` and the number outside the sub-field search range.

  Separately, "DOI:" prefix lines (no annotated tokens) were causing the label function to return
  `None`, which backed the TEI writer up to `<listBibl>` level and produced a spurious `<bibl>` for
  the URL on the next line. Lines with no annotated reference tokens now continue the current bibl
  as `<reference>` when a reference is already open.
---
 .../training/cli/generate_data.py             |   5 +-
 sciencebeam_parser/training/jats/aligner.py   | 100 ++++++++++++++++++
 tests/training/cli/generate_data_test.py      |  31 ++++++
 tests/training/jats/test_aligner.py           |  59 +++++++++++
 4 files changed, 194 insertions(+), 1 deletion(-)

diff --git a/sciencebeam_parser/training/cli/generate_data.py b/sciencebeam_parser/training/cli/generate_data.py
index b8510816..e5488cec 100644
--- a/sciencebeam_parser/training/cli/generate_data.py
+++ b/sciencebeam_parser/training/cli/generate_data.py
@@ -1131,7 +1131,10 @@ def fn(
                 )
             line_instance_id = line_first_instance_cache[line_id]
             if line_instance_id is None:
-                return None
+                # Gap-fill: a line with no annotated reference tokens (e.g. "DOI:"
+                # prefix before the URL) continues the current bibl rather than
+                # breaking out to listBibl level.
+                return '<reference>' if last_instance_id is not None else None
             if line_instance_id != last_instance_id and last_instance_id is not None:
                 last_instance_id = line_instance_id
                 return 'B-<reference>'
diff --git a/sciencebeam_parser/training/jats/aligner.py b/sciencebeam_parser/training/jats/aligner.py
index bf2bd211..97354b72 100644
--- a/sciencebeam_parser/training/jats/aligner.py
+++ b/sciencebeam_parser/training/jats/aligner.py
@@ -1,3 +1,4 @@
+# pylint: disable=too-many-lines
 import logging
 import re
 from dataclasses import dataclass
@@ -363,6 +364,100 @@ def _is_pure_number(text: str) -> bool:
     return bool(re.fullmatch(r'\d+', text))
 
 
+_BRACKET_LABEL_RE = re.compile(r'^(\[)(.+)(\])$|^(\()(.+)(\))$')
+
+# How far before the first segment start to search for the bracket label inner content.
+# Needed because the parent SW match may start at "]" (skipping the preceding "[" and
+# the inner number), so the sub-field search range starts after the label tokens.
+_BRACKET_LABEL_BACK_BUFFER = 10
+
+
+def _find_inner_token(
+    token_index: _TokenIndex,
+    inner: str,
+    segments: List[Tuple[int, int]],
+) -> Optional[_MatchResult]:
+    """Locate `inner` as a standalone token within segments; pure-number fast path."""
+    if _is_pure_number(inner):
+        return _exact_number_match(token_index, inner, segments)
+    haystack = token_index.haystack
+    inner_len = len(inner)
+    for seg_start, seg_end in segments:
+        pos = seg_start
+        while pos <= seg_end - inner_len:
+            idx = haystack.find(inner, pos, seg_end)
+            if idx == -1:
+                break
+            end = idx + inner_len
+            if (token_index.is_in_token(idx)
+                    and token_index.is_token_start(idx)
+                    and token_index.is_token_boundary_after(end - 1)):
+                return (idx, end, [(idx, end)])
+            pos = idx + 1
+    return None
+
+
+def _try_bracket_label_match(  # pylint: disable=too-many-locals
+    token_index: _TokenIndex,
+    needle: str,
+    segments: List[Tuple[int, int]],
+) -> Optional[_MatchResult]:
+    """Match bracket-style labels like [1] or (2) whose tokens are split by the PDF tokeniser.
+
+    When the PDF tokeniser produces three tokens "[", "1", "]" the haystack has
+    "[ 1 ]".  Smith-Waterman cannot match "[1]" across those spaces because the
+    scoring library's traceback terminates early at gap moves, yielding quality < threshold.
+
+    This function strips the outer brackets, matches the inner content, then extends
+    the match range to cover the bracket tokens that immediately surround the hit.
+    The search extends _BRACKET_LABEL_BACK_BUFFER chars before the first segment start
+    because the parent SW match often begins at "]", leaving "[" and the number before
+    its first matched block (and thus outside the nominal search range).
+    """
+    m = _BRACKET_LABEL_RE.fullmatch(needle)
+    if not m:
+        return None
+    if m.group(1):
+        open_b, inner, close_b = m.group(1), m.group(2), m.group(3)
+    else:
+        open_b, inner, close_b = m.group(4), m.group(5), m.group(6)
+
+    extended = (
+        [(max(0, segments[0][0] - _BRACKET_LABEL_BACK_BUFFER), segments[0][1])]
+        + list(segments[1:])
+    ) if segments else segments
+
+    inner_match = _find_inner_token(token_index, inner, extended)
+    if inner_match is None:
+        return None
+
+    m_start, m_end, blocks = inner_match
+    haystack = token_index.haystack
+
+    # Extend to include adjacent opening bracket token (up to 2 chars before m_start).
+    new_start = m_start
+    new_blocks: List[Tuple[int, int]] = list(blocks)
+    for pos in range(max(0, m_start - 2), m_start):
+        if (haystack[pos] == open_b
+                and token_index.is_in_token(pos)
+                and token_index.is_token_start(pos)):
+            new_start = pos
+            new_blocks = [(pos, pos + 1)] + new_blocks
+            break
+
+    # Extend to include adjacent closing bracket token (up to 2 chars after m_end).
+    new_end = m_end
+    for pos in range(m_end, min(m_end + 2, len(haystack))):
+        if (haystack[pos] == close_b
+                and token_index.is_in_token(pos)
+                and token_index.is_token_start(pos)):
+            new_end = pos + 1
+            new_blocks = new_blocks + [(pos, pos + 1)]
+            break
+
+    return new_start, new_end, new_blocks
+
+
 def _is_exact_sw_match(result: _MatchResult, needle_len: int) -> bool:
     """True when SW found the needle as one contiguous block (no gaps)."""
     _, _, blocks = result
@@ -437,6 +532,11 @@ def _fuzzy_match_field_value(  # pylint: disable=too-many-locals
                 break
             start += stride
 
+    if gap_match is None:
+        bracket_match = _try_bracket_label_match(token_index, needle, segments)
+        if bracket_match is not None:
+            return bracket_match
+
     return gap_match
 
 
diff --git a/tests/training/cli/generate_data_test.py b/tests/training/cli/generate_data_test.py
index d531d743..64b1cacb 100644
--- a/tests/training/cli/generate_data_test.py
+++ b/tests/training/cli/generate_data_test.py
@@ -1019,6 +1019,37 @@ def test_label_token_at_instance_transition_emits_b_label(self):
         for idx in range(1, len(line2.tokens)):
             assert label_fn(annotated, {}, _make_md(line2, idx)) == '<reference>'
 
+    def test_unannotated_line_between_references_is_gap_filled(self):
+        # "DOI:" may appear on a line by itself with no annotated reference tokens
+        # (JATS only stores the DOI value, not the "DOI:" prefix).  Returning None
+        # for that line would back the TEI writer up to listBibl level, creating a
+        # spurious <bibl> for the URL that follows.  The fix gap-fills to '<reference>'
+        # so "DOI:" stays inside the current bibl.
+        line1 = LayoutLine.for_text('Smith J 2020 A title DOI')
+        line_doi = LayoutLine.for_text('DOI:')          # no annotated tokens
+        line_url = LayoutLine.for_text('10.1234/abcd')
+        refs_doc = LayoutDocument(pages=[LayoutPage(blocks=[
+            LayoutBlock(lines=[line1, line_doi, line_url])
+        ])])
+        annotated = JatsAnnotatedLayoutDocument(layout_document=refs_doc)
+        for t in line1.tokens:
+            annotated.set_token_label(t, JatsFieldNames.REFERENCE, instance_id=1)
+        for t in line_url.tokens:
+            annotated.set_token_label(t, JatsFieldNames.REFERENCE, instance_id=1)
+        # line_doi tokens are intentionally unannotated
+
+        label_fn = ReferenceSegmenterModelTrainingDataGenerator().get_jats_label_fn()
+        assert label_fn is not None
+
+        # Advance through line1 tokens
+        for idx in range(len(line1.tokens)):
+            label_fn(annotated, {}, _make_md(line1, idx))
+        # "DOI:" line has no annotations → gap-fill returns '<reference>', not None
+        doi_label = label_fn(annotated, {}, _make_md(line_doi, 0))
+        assert doi_label == '<reference>', (
+            f'Unannotated line after a reference should gap-fill to <reference>, got {doi_label!r}'
+        )
+
 
 # ── CitationModelTrainingDataGenerator JATS path ──────────────────────────────
 
diff --git a/tests/training/jats/test_aligner.py b/tests/training/jats/test_aligner.py
index af6e0cb4..af95fcb0 100644
--- a/tests/training/jats/test_aligner.py
+++ b/tests/training/jats/test_aligner.py
@@ -628,6 +628,65 @@ def test_regular_fallback_extends_with_given_names(self):
             f'"Y" initial should be labeled REFERENCE_AUTHOR after fallback, got {sub!r}'
         )
 
+    def test_bracket_label_tokens_are_labeled_reference_label(self):
+        # Reference labels like "[1]" tokenize as three separate tokens "[", "1", "]".
+        # The haystack has "[ 1 ]" with spaces; SW cannot match "[1]" contiguously.
+        # _try_bracket_label_match strips the brackets, finds "1" via exact match,
+        # then extends the range to include the adjacent bracket tokens.
+        doc = _make_doc(
+            '[ 1 ] Richards FA A flexible growth function 1959',
+            '[ 2 ] Jones B Another title 2020',
+        )
+        fvs = [
+            _fv('[1] Richards FA A flexible growth function 1959', JatsFieldNames.REFERENCE),
+            _fv('[1]', JatsFieldNames.REFERENCE, JatsSubFieldNames.REFERENCE_LABEL),
+            _fv('Richards FA', JatsFieldNames.REFERENCE, JatsSubFieldNames.REFERENCE_AUTHOR),
+            _fv('1959', JatsFieldNames.REFERENCE, JatsSubFieldNames.REFERENCE_YEAR),
+            _fv('[2] Jones B Another title 2020', JatsFieldNames.REFERENCE),
+            _fv('[2]', JatsFieldNames.REFERENCE, JatsSubFieldNames.REFERENCE_LABEL),
+            _fv('Jones B', JatsFieldNames.REFERENCE, JatsSubFieldNames.REFERENCE_AUTHOR),
+            _fv('2020', JatsFieldNames.REFERENCE, JatsSubFieldNames.REFERENCE_YEAR),
+        ]
+        annotated = self._align(doc, fvs)
+        tokens = list(doc.iter_all_tokens())
+        label_tokens = [
+            t for t in tokens
+            if annotated.get_token_sub_field(t) == JatsSubFieldNames.REFERENCE_LABEL
+        ]
+        label_texts = [t.text for t in label_tokens]
+        assert label_texts == ['[', '1', ']', '[', '2', ']'], (
+            f'Expected 6 bracket label tokens, got {label_texts!r}'
+        )
+
+    def test_bracket_label_tokens_labeled_when_parent_starts_at_closing_bracket(self):
+        # The parent SW match for "[1] Richards..." often starts at "]" (position 4 in
+        # "[ 1 ] ...") because the SW can't match "[1]" across spaces and skips to "]".
+        # The REFERENCE_LABEL sub-field search must extend backward to find "[" and "1".
+        doc = _make_doc('[ 1 ] Richards FA long title here Journal 1999')
+        fvs = [
+            _fv(
+                '[1] Richards FA long title here Journal 1999',
+                JatsFieldNames.REFERENCE,
+            ),
+            _fv('[1]', JatsFieldNames.REFERENCE, JatsSubFieldNames.REFERENCE_LABEL),
+            _fv('long title here', JatsFieldNames.REFERENCE,
+                JatsSubFieldNames.REFERENCE_ARTICLE_TITLE),
+            _fv('Journal', JatsFieldNames.REFERENCE, JatsSubFieldNames.REFERENCE_SOURCE),
+            _fv('1999', JatsFieldNames.REFERENCE, JatsSubFieldNames.REFERENCE_YEAR),
+        ]
+        annotated = self._align(doc, fvs)
+        tokens = list(doc.iter_all_tokens())
+        bracket_open = next((t for t in tokens if t.text == '['), None)
+        assert bracket_open is not None
+        assert annotated.get_token_sub_field(bracket_open) == JatsSubFieldNames.REFERENCE_LABEL, (
+            '"[" should be labeled REFERENCE_LABEL, not left unlabeled'
+        )
+        num_token = next((t for t in tokens if t.text == '1'), None)
+        assert num_token is not None
+        assert annotated.get_token_sub_field(num_token) == JatsSubFieldNames.REFERENCE_LABEL, (
+            '"1" should be labeled REFERENCE_LABEL'
+        )
+
     def test_masked_sub_field_prevents_duplicate_match(self):
         # When the same author name appears twice in a reference, masking ensures
         # the second sub-field value matches the second occurrence rather than

From abd3289db392bf07854aabae8999760dd9da43d4 Mon Sep 17 00:00:00 2001
From: Daniel Ecer <de-code@users.noreply.github.com>
Date: Fri, 26 Jun 2026 21:22:36 +0100
Subject: [PATCH 04/12] Fix reference parent match quality to recover missing
 scielo bibls
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

  Two fixes reduce the scielo citation score regression introduced when
  JATS-guided bibl splitting replaced ML-model splitting:

  1. field_extractor.py: strip <pub-id> subtrees from the parent
     reference text used for Smith-Waterman alignment.  PMID/PMCID values
     appear in JATS but not in PDF reference lists; including them pads the
     needle length without adding matched characters, pushing quality below
     the 0.8 threshold.

  2. aligner.py: when a REFERENCE parent match fails at 0.8, retry with a
     relaxed 0.65 threshold.  JATS may concatenate author initials ("CA")
     or order publisher/place differently from the PDF, producing valid
     near-threshold alignments that the primary pass rejects.

  Both issues caused parent matches to silently fail, skipping a bibl
  entirely.  Because partial_list scoring is ordered, one skipped bibl
  shifts every subsequent field comparison by one position, cascading into
  many false mismatches (e.g. PPR489375 dropped from 22/22 to 3/22 correct
  issue matches due to a single missing bibl at position 2).

  After these fixes the scielo regression shrinks from −0.124 to −0.009
  mean while ORE remains at +0.363 mean, for a net +0.177 on citation.
---
 sciencebeam_parser/training/jats/aligner.py   | 28 +++++++++++++++++++
 .../training/jats/field_extractor.py          | 27 +++++++++++++++++-
 2 files changed, 54 insertions(+), 1 deletion(-)

diff --git a/sciencebeam_parser/training/jats/aligner.py b/sciencebeam_parser/training/jats/aligner.py
index 97354b72..3b071eba 100644
--- a/sciencebeam_parser/training/jats/aligner.py
+++ b/sciencebeam_parser/training/jats/aligner.py
@@ -118,6 +118,13 @@
 #   (abs_start, abs_end, [(block_start, block_end), ...])
 _MatchResult = Tuple[int, int, List[Tuple[int, int]]]
 
+# When a parent REFERENCE match fails at the primary threshold (0.8), retry at this
+# lower value.  JATS author initials may be concatenated ("CA") while the PDF expands
+# them ("C. A."), and institutional refs can omit boilerplate text that pads the needle
+# without appearing in the PDF reference list.  Keeping this above 0.5 is enough to
+# reject genuinely absent references while recovering these near-threshold cases.
+_REFERENCE_PARENT_MIN_THRESHOLD = 0.65
+
 
 @dataclass
 class AlignmentConfig:
@@ -967,6 +974,27 @@ def align(  # pylint: disable=too-many-locals,too-many-branches,too-many-stateme
                     token_index, fv, self.config,
                     search_start=body_floor, search_end=None,
                 )
+            # Parent REFERENCE fallback: retry with a relaxed threshold when the
+            # full-text parent match just misses 0.8.  JATS may concatenate initials
+            # ("CA") or order publisher/place differently from the PDF reference list,
+            # reducing quality without indicating a wrong match.  Only applied to
+            # parent matches (sub_field_name is None) of the REFERENCE field so that
+            # sub-field containment and other fields keep the stricter threshold.
+            if (
+                match_range is None
+                and fv.sub_field_name is None
+                and fv.field_name in _REFERENCE_FIELDS
+                and _REFERENCE_PARENT_MIN_THRESHOLD < self.config.threshold
+            ):
+                _relaxed_config = AlignmentConfig(
+                    threshold=_REFERENCE_PARENT_MIN_THRESHOLD,
+                    max_window=self.config.max_window,
+                )
+                match_range = _fuzzy_match_field_value(
+                    token_index, fv, _relaxed_config,
+                    search_start=search_start, search_end=search_end,
+                    masked_ranges=masked,
+                )
             # Sub-field fallback: retry with fallback_text (e.g. surname only)
             # when the primary JATS name text does not match the PDF text.
             if match_range is None and fv.sub_field_name is not None and fv.fallback_text:
diff --git a/sciencebeam_parser/training/jats/field_extractor.py b/sciencebeam_parser/training/jats/field_extractor.py
index fea173b8..e9596330 100644
--- a/sciencebeam_parser/training/jats/field_extractor.py
+++ b/sciencebeam_parser/training/jats/field_extractor.py
@@ -21,6 +21,31 @@ def _element_text(el: etree._Element) -> str:
     return ' '.join(' '.join(el.itertext()).split())
 
 
+def _reference_parent_text(ref_el: etree._Element) -> str:
+    """Extract parent-match text for a <ref>, skipping <pub-id> subtrees.
+
+    PMID/PMCID values appear in JATS <pub-id> elements but rarely in PDF
+    reference lists.  Including them in the SW needle inflates its length
+    without adding matched characters, pushing match quality below the 0.8
+    threshold and causing the entire reference to go unmatched.  DOI text is
+    also excluded here — it is handled by the dedicated REFERENCE_DOI
+    sub-field match, so the parent only needs the bibliographic prose.
+    """
+    parts: List[str] = []
+    for child in ref_el.iter():
+        tag = child.tag
+        local = tag.split('}', 1)[1] if isinstance(tag, str) and tag.startswith('{') else tag
+        if local == 'pub-id':
+            if child.tail:
+                parts.append(child.tail)
+            continue
+        if child.text:
+            parts.append(child.text)
+        if child is not ref_el and child.tail:
+            parts.append(child.tail)
+    return ' '.join(' '.join(parts).split())
+
+
 def _iter_sub_field_values(
     parent_el: etree._Element,
     field_name: str,
@@ -408,7 +433,7 @@ def _iter_back_reference_values(self, root: etree._Element) -> Iterator[JatsFiel
                 )
 
         for ref_el in root.xpath('back/ref-list/ref'):
-            text = _element_text(ref_el)
+            text = _reference_parent_text(ref_el)
             if text:
                 yield JatsFieldValue(text=text, field_name=JatsFieldNames.REFERENCE)
             yield from _iter_reference_author_values(ref_el, JatsFieldNames.REFERENCE)

From bf45954c9e67f33f3e7b81db8361ac8ea47e2251 Mon Sep 17 00:00:00 2001
From: Daniel Ecer <d.ecer@elifesciences.org>
Date: Mon, 29 Jun 2026 09:15:32 +0100
Subject: [PATCH 05/12] Fix chapter-title extraction and page range attributes
 in citation training data

- Add <chapter-title> to REFERENCE_ARTICLE_TITLE XPath so book-chapter and
  government-document refs are annotated with titles (recovers ~31 missing
  titles across scielo_preprints-jats documents)
- Post-process <biblScope unit="page"> elements to add @from/@to attributes
  when the element contains a fpage-lpage range ("388 - 412"), matching the
  format expected by sciencebeam-judge for correct fpage extraction

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
---
 sciencebeam_parser/models/citation/training_data.py | 11 +++++++++++
 sciencebeam_parser/training/jats/field_extractor.py |  2 +-
 2 files changed, 12 insertions(+), 1 deletion(-)

diff --git a/sciencebeam_parser/models/citation/training_data.py b/sciencebeam_parser/models/citation/training_data.py
index 649ea01d..65210abe 100644
--- a/sciencebeam_parser/models/citation/training_data.py
+++ b/sciencebeam_parser/models/citation/training_data.py
@@ -1,4 +1,5 @@
 import logging
+import re
 
 from lxml import etree
 
@@ -12,6 +13,9 @@
 )
 from sciencebeam_parser.utils.xml import get_text_content
 
+# Matches "388 - 412", "281-282", "1199 -1207" etc.
+_PAGE_RANGE_RE = re.compile(r'^(\S+)\s*[-–]\s*(\S+)$')
+
 
 LOGGER = logging.getLogger(__name__)
 
@@ -68,6 +72,13 @@ def get_post_processed_xml_root(self, xml_root: etree.ElementBase):
             if not external_identifier_type:
                 continue
             idno_element.attrib['type'] = external_identifier_type
+        for page_el in tei_xpath(xml_root, '//tei:biblScope[@unit="page"]'):
+            full_text = ' '.join(page_el.itertext()).replace('\n', ' ').strip()
+            full_text = re.sub(r'\s+', ' ', full_text)
+            m = _PAGE_RANGE_RE.match(full_text)
+            if m and 'from' not in page_el.attrib:
+                page_el.attrib['from'] = m.group(1)
+                page_el.attrib['to'] = m.group(2)
         return xml_root
 
 
diff --git a/sciencebeam_parser/training/jats/field_extractor.py b/sciencebeam_parser/training/jats/field_extractor.py
index e9596330..5f4430b6 100644
--- a/sciencebeam_parser/training/jats/field_extractor.py
+++ b/sciencebeam_parser/training/jats/field_extractor.py
@@ -66,7 +66,7 @@ def _iter_sub_field_values(
 # Sub-field XPaths for references (relative to each <ref> element)
 _REFERENCE_SUB_FIELDS = [
     (JatsSubFieldNames.REFERENCE_LABEL,           './label'),
-    (JatsSubFieldNames.REFERENCE_ARTICLE_TITLE,   './/article-title'),
+    (JatsSubFieldNames.REFERENCE_ARTICLE_TITLE,   './/article-title | .//chapter-title'),
     (JatsSubFieldNames.REFERENCE_SOURCE,          './/source'),
     (JatsSubFieldNames.REFERENCE_YEAR,            './/year'),
     (JatsSubFieldNames.REFERENCE_VOLUME,          './/volume'),

From 55f720c8b5928d5e4b01301d2147ac67b0c3c97c Mon Sep 17 00:00:00 2001
From: Daniel Ecer <d.ecer@elifesciences.org>
Date: Mon, 29 Jun 2026 10:35:19 +0100
Subject: [PATCH 06/12] Lower REFERENCE parent match fallback threshold from
 0.65 to 0.55
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

A truncated PDF reference (one where year/volume/URL are absent from the
scanned text) produces a JATS needle much longer than the matching haystack
fragment, giving quality ≈ 0.60.  The previous 0.65 floor rejected these,
leaving the bibl unmatched and cascading into position-shifted comparisons
for every subsequent ordered field (reference_issue, reference_volume).

Lowering to 0.55 captures these cases while still rejecting genuinely
absent references.  When the parent match now succeeds, the aligner can
search for sub-fields beyond the reference-segmenter bibl boundary —
recovering the source annotation and any continuation text the segmenter
missed — producing a complete, separate citation bibl for the reference.

Scielo impact (PPR459390 case):
  reference_issue/exact:   0.943 → 0.987 (+0.044)
  reference_volume/exact:  0.979 → 0.998 (+0.019)
  reference_source/lev:    0.984 → 0.996 (+0.012)
  scielo mean:             0.908 → 0.919 (+0.011)

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
---
 sciencebeam_parser/training/jats/aligner.py | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/sciencebeam_parser/training/jats/aligner.py b/sciencebeam_parser/training/jats/aligner.py
index 3b071eba..d058e789 100644
--- a/sciencebeam_parser/training/jats/aligner.py
+++ b/sciencebeam_parser/training/jats/aligner.py
@@ -121,9 +121,10 @@
 # When a parent REFERENCE match fails at the primary threshold (0.8), retry at this
 # lower value.  JATS author initials may be concatenated ("CA") while the PDF expands
 # them ("C. A."), and institutional refs can omit boilerplate text that pads the needle
-# without appearing in the PDF reference list.  Keeping this above 0.5 is enough to
-# reject genuinely absent references while recovering these near-threshold cases.
-_REFERENCE_PARENT_MIN_THRESHOLD = 0.65
+# without appearing in the PDF reference list.  0.55 is sufficient to capture truncated
+# PDF references (e.g. a ref whose year/volume/URL are absent, giving quality ~0.60)
+# while rejecting genuinely absent ones.
+_REFERENCE_PARENT_MIN_THRESHOLD = 0.55
 
 
 @dataclass

From f5c8af63a3f1341c7b57e3578e2155146c1d3304 Mon Sep 17 00:00:00 2001
From: Daniel Ecer <d.ecer@elifesciences.org>
Date: Mon, 29 Jun 2026 11:47:38 +0100
Subject: [PATCH 07/12] Fix single-digit reference labels not annotated when
 label precedes parent match
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Labels "1."–"9." appear immediately before the JATS parent text in the PDF (the
parent text excludes the label), so the Smith-Waterman parent match starts after
them.  The sub-field search range began at p_start, placing single-digit labels
outside it.

Two-part fix:
- Wire in _SUB_FIELD_PARENT_PRE_BUFFER (was defined but unused) and set it to 20
  chars.  Applied only for REFERENCE_LABEL sub-fields so other sub-fields (e.g.
  country) cannot false-match an identical word in an earlier context.
- Add _is_punct_suffix_token helper so "1." tokenised as a single token passes
  _exact_number_match even though is_token_boundary_after fails for the digit alone.

Improves citation label annotation (labels 1-9 now appear as <note>N</note>) and
reference-segmenter label score (0.870 → 0.964 on scielo_preprints-jats).

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
---
 sciencebeam_parser/training/jats/aligner.py | 39 +++++++++++++++++++--
 1 file changed, 36 insertions(+), 3 deletions(-)

diff --git a/sciencebeam_parser/training/jats/aligner.py b/sciencebeam_parser/training/jats/aligner.py
index d058e789..5d0c3e0e 100644
--- a/sciencebeam_parser/training/jats/aligner.py
+++ b/sciencebeam_parser/training/jats/aligner.py
@@ -88,7 +88,10 @@
 # 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
+# How far before the parent match start to extend the sub-field search.  Labels like
+# "1." appear immediately before the parent text (which is JATS without the label), so
+# the parent SW match starts after them.  A small back-buffer re-includes them.
+_SUB_FIELD_PARENT_PRE_BUFFER = 20
 
 # Anchor+chain labelling strategy:
 # Smith-Waterman produces many tiny (1–4 char) matching blocks while traversing
@@ -472,6 +475,26 @@ def _is_exact_sw_match(result: _MatchResult, needle_len: int) -> bool:
     return len(blocks) == 1 and (blocks[0][1] - blocks[0][0]) == needle_len
 
 
+def _is_punct_suffix_token(
+    token_index: _TokenIndex,
+    haystack: str,
+    end: int,
+) -> bool:
+    """Return True when the characters after `end` (still in the same token) are all
+    punctuation.  This covers "1." or "1," where the PDF tokeniser attaches the
+    delimiter to the digit, so the exact number "1" cannot be matched with a clean
+    token boundary but is still the correct label to extract."""
+    pos = end
+    while pos < len(haystack) and token_index.is_in_token(pos):
+        if not haystack[pos].isspace() and haystack[pos] not in '.,;:)]}':
+            return False
+        if not token_index.is_token_boundary_after(pos):
+            pos += 1
+            continue
+        break
+    return True
+
+
 def _exact_number_match(
     token_index: _TokenIndex,
     needle: str,
@@ -488,7 +511,8 @@ def _exact_number_match(
             end = idx + needle_len
             if (token_index.is_in_token(idx)
                     and token_index.is_token_start(idx)
-                    and token_index.is_token_boundary_after(end - 1)):
+                    and (token_index.is_token_boundary_after(end - 1)
+                         or _is_punct_suffix_token(token_index, haystack, end))):
                 return idx, end, [(idx, end)]
             pos = idx + 1
     return None
@@ -561,7 +585,16 @@ def _search_range(
     """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
+        # Extend backward only for labels: the JATS parent text excludes the label
+        # ("1.", "2.", ...), so the parent match starts after it.  Other sub-fields
+        # (e.g. country) must stay strictly inside the parent range to avoid matching
+        # identical text in an earlier reference or affiliation.
+        pre = (
+            _SUB_FIELD_PARENT_PRE_BUFFER
+            if fv.sub_field_name == JatsSubFieldNames.REFERENCE_LABEL
+            else 0
+        )
+        return max(0, p_start - pre), 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:

From f2550774aeac6117f832dea127481dab4e0d834f Mon Sep 17 00:00:00 2001
From: Daniel Ecer <d.ecer@elifesciences.org>
Date: Mon, 29 Jun 2026 13:42:41 +0100
Subject: [PATCH 08/12] Restrict reference label pre-buffer to pure-number
 labels only
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

  REFERENCE_LABEL sub-fields with a suffix (e.g. "20-" in scielo_preprints
  references like "1-", "20-") are included in the JATS parent reference text,
  so p_start is already at the label — extending the search backward would let
  "20-" false-match the "20-" fragment inside a preceding DOI like "2020-0248".

  The fix applies the pre-buffer only when the label text is a pure number
  (digits only). Suffixed labels ("1-", "20-") and all other sub-fields keep
  pre=0 and are confined to the parent match range.

  The pre-buffer value stays at 20 to handle cases where the JATS parent text
  starts after a publisher name ("Brasil."), leaving a ~14-char gap between the
  label token and p_start.

  Adds three unit tests covering the positive case, the DOI false-match
  regression, and the long-gap (extra publisher text) scenario.
---
 sciencebeam_parser/training/jats/aligner.py | 16 ++---
 tests/training/jats/test_aligner.py         | 72 +++++++++++++++++++++
 2 files changed, 77 insertions(+), 11 deletions(-)

diff --git a/sciencebeam_parser/training/jats/aligner.py b/sciencebeam_parser/training/jats/aligner.py
index 5d0c3e0e..98e60ecf 100644
--- a/sciencebeam_parser/training/jats/aligner.py
+++ b/sciencebeam_parser/training/jats/aligner.py
@@ -88,9 +88,7 @@
 # matched range.  Keeps short sub-field values (e.g. "USA", "2020") from matching
 # identical text elsewhere in the document.
 _SUB_FIELD_PARENT_BUFFER = 200
-# How far before the parent match start to extend the sub-field search.  Labels like
-# "1." appear immediately before the parent text (which is JATS without the label), so
-# the parent SW match starts after them.  A small back-buffer re-includes them.
+# Pure-number labels precede the JATS parent text; this extends the search backward.
 _SUB_FIELD_PARENT_PRE_BUFFER = 20
 
 # Anchor+chain labelling strategy:
@@ -585,15 +583,11 @@ def _search_range(
     """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]
-        # Extend backward only for labels: the JATS parent text excludes the label
-        # ("1.", "2.", ...), so the parent match starts after it.  Other sub-fields
-        # (e.g. country) must stay strictly inside the parent range to avoid matching
-        # identical text in an earlier reference or affiliation.
-        pre = (
-            _SUB_FIELD_PARENT_PRE_BUFFER
-            if fv.sub_field_name == JatsSubFieldNames.REFERENCE_LABEL
-            else 0
+        is_pure_number_label = (
+            fv.sub_field_name == JatsSubFieldNames.REFERENCE_LABEL
+            and _is_pure_number(fv.text)
         )
+        pre = _SUB_FIELD_PARENT_PRE_BUFFER if is_pure_number_label else 0
         return max(0, p_start - pre), 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
diff --git a/tests/training/jats/test_aligner.py b/tests/training/jats/test_aligner.py
index af95fcb0..aa7600ca 100644
--- a/tests/training/jats/test_aligner.py
+++ b/tests/training/jats/test_aligner.py
@@ -1,3 +1,4 @@
+# pylint: disable=too-many-lines
 from typing import Optional
 
 from sciencebeam_parser.document.layout_document import (
@@ -973,3 +974,74 @@ def test_hyphenated_surname_labeled_when_jats_has_extra_given_name(self):
         assert '-' in author_tokens, (
             f'Hyphen in "Silva-Braun" must be REFERENCE_AUTHOR; got: {author_tokens}'
         )
+
+    def test_pure_number_label_found_before_parent_match(self):
+        # JATS strips the numeric label from the parent reference text, so the SW
+        # match for the parent starts after the label token.  The pre-buffer extends the
+        # sub-field search just far enough back to reach "1" before "Smith".
+        doc = _make_doc('1. Smith J title here 2020')
+        fvs = [
+            _fv('Smith J title here 2020', JatsFieldNames.REFERENCE),
+            _fv('1', JatsFieldNames.REFERENCE, JatsSubFieldNames.REFERENCE_LABEL),
+            _fv('Smith J', JatsFieldNames.REFERENCE, JatsSubFieldNames.REFERENCE_AUTHOR),
+            _fv('2020', JatsFieldNames.REFERENCE, JatsSubFieldNames.REFERENCE_YEAR),
+        ]
+        annotated = self._align(doc, fvs)
+        tokens = list(doc.iter_all_tokens())
+        label_token = next((t for t in tokens if t.text == '1'), None)
+        assert label_token is not None, 'Expected "1" token in doc'
+        sub = annotated.get_token_sub_field(label_token)
+        assert sub == JatsSubFieldNames.REFERENCE_LABEL, (
+            f'"1" should be REFERENCE_LABEL; got {sub!r}'
+        )
+
+    def test_suffixed_label_not_false_matched_in_preceding_doi(self):
+        # Suffixed labels like "20-" are not pure numbers, so no pre-buffer is applied.
+        # Without this guard, "20" inside the parent text would also match the "20" inside
+        # a preceding DOI fragment like "2020-0248".
+        doc = _make_doc(
+            'Preceding ref doi 2020-0248',
+            '20- Kato T study title Stress Health 2015',
+        )
+        fvs = [
+            _fv('Preceding ref doi 2020-0248', JatsFieldNames.REFERENCE),
+            _fv('20- Kato T study title Stress Health 2015', JatsFieldNames.REFERENCE),
+            _fv('20-', JatsFieldNames.REFERENCE, JatsSubFieldNames.REFERENCE_LABEL),
+            _fv('Kato T', JatsFieldNames.REFERENCE, JatsSubFieldNames.REFERENCE_AUTHOR),
+        ]
+        annotated = self._align(doc, fvs)
+        tokens = list(doc.iter_all_tokens())
+        # "2020" from the first line must not be labeled as REFERENCE_LABEL
+        doi_2020 = next((t for t in tokens if t.text == '2020'), None)
+        assert doi_2020 is not None
+        assert annotated.get_token_sub_field(doi_2020) != JatsSubFieldNames.REFERENCE_LABEL, (
+            '"2020" in the preceding DOI must not be labeled REFERENCE_LABEL'
+        )
+        # "20" from the second line (the actual label start) must be labeled
+        label_tokens = [
+            t for t in tokens
+            if annotated.get_token_sub_field(t) == JatsSubFieldNames.REFERENCE_LABEL
+        ]
+        label_texts = [t.text for t in label_tokens]
+        assert '20' in label_texts, (
+            f'"20" on line 2 should be REFERENCE_LABEL; got {label_texts!r}'
+        )
+
+    def test_pure_number_label_found_with_extra_text_before_parent(self):
+        # When the JATS parent text begins at the publication name (e.g. "Emenda
+        # Constitucional..."), the label and the publisher name ("Brasil.") both sit
+        # before p_start — the gap can be ~14 chars.  The pre-buffer must cover it.
+        doc = _make_doc('17. Brasil. Title about health 2020')
+        fvs = [
+            _fv('Title about health 2020', JatsFieldNames.REFERENCE),
+            _fv('17', JatsFieldNames.REFERENCE, JatsSubFieldNames.REFERENCE_LABEL),
+            _fv('2020', JatsFieldNames.REFERENCE, JatsSubFieldNames.REFERENCE_YEAR),
+        ]
+        annotated = self._align(doc, fvs)
+        tokens = list(doc.iter_all_tokens())
+        label_token = next((t for t in tokens if t.text == '17'), None)
+        assert label_token is not None, 'Expected "17" token in doc'
+        sub = annotated.get_token_sub_field(label_token)
+        assert sub == JatsSubFieldNames.REFERENCE_LABEL, (
+            f'"17" should be REFERENCE_LABEL even when parent starts 14 chars later; got {sub!r}'
+        )

From af1215ceb7dec00e6f0ed74accfc21b41a715813 Mon Sep 17 00:00:00 2001
From: Daniel Ecer <d.ecer@elifesciences.org>
Date: Mon, 29 Jun 2026 14:24:58 +0100
Subject: [PATCH 09/12] Fix single-digit suffixed reference labels (e.g. "1-")
 not being annotated
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

  When a PDF tokenises "1-" as two tokens "1" and "-", Smith-Waterman
  cannot match the 2-char needle "1-" against the spaced haystack "1 - …"
  with sufficient quality (matched=1/2=0.5 < threshold=0.8).  The parent
  SW match compounds the issue: it aligns the full reference text by
  skipping the leading digit and starting at "-" (p_start=2), so the
  label sub-field search range excluded position 0 even for pure-digit
  prefixes of suffixed labels.

  Two targeted fixes in aligner.py:

  1. _search_range: apply a small pre-buffer (_SUB_FIELD_LABEL_DIGIT_PRE_BUFFER=3)
     for any REFERENCE_LABEL whose text starts with digits followed by a
     suffix (e.g. "1-", "9-").  This shifts search_start back to 0 so the
     digit token is within range.

  2. _fuzzy_match_field_value: add _try_numeric_prefix_label_match as a
     final fallback for REFERENCE_LABEL sub-fields when SW fails.  It
     splits the label into a digit prefix and suffix, locates the prefix
     via _exact_number_match (which enforces token-start boundaries and
     therefore cannot false-match inside "2020" or "0248"), then extends
     the match to cover the immediately adjacent suffix token(s).

  Two-digit suffixed labels (e.g. "10-") continue to be found via the
  normal SW path (quality=1.0) and are unaffected.  The existing
  regression guard for "20-" not matching "2020-0248" is preserved.
---
 sciencebeam_parser/training/jats/aligner.py | 45 ++++++++++++++++++++-
 tests/training/jats/test_aligner.py         | 36 +++++++++++++++++
 2 files changed, 80 insertions(+), 1 deletion(-)

diff --git a/sciencebeam_parser/training/jats/aligner.py b/sciencebeam_parser/training/jats/aligner.py
index 98e60ecf..2ef05b76 100644
--- a/sciencebeam_parser/training/jats/aligner.py
+++ b/sciencebeam_parser/training/jats/aligner.py
@@ -90,6 +90,10 @@
 _SUB_FIELD_PARENT_BUFFER = 200
 # Pure-number labels precede the JATS parent text; this extends the search backward.
 _SUB_FIELD_PARENT_PRE_BUFFER = 20
+# Digit-prefix labels (e.g. "1-") may also precede the parent match start because SW
+# aligns the parent by skipping the label digit, leaving p_start at the suffix char.
+# A small buffer is enough — single-digit + space = 2 chars.
+_SUB_FIELD_LABEL_DIGIT_PRE_BUFFER = 3
 
 # Anchor+chain labelling strategy:
 # Smith-Waterman produces many tiny (1–4 char) matching blocks while traversing
@@ -374,6 +378,7 @@ def _is_pure_number(text: str) -> bool:
 
 
 _BRACKET_LABEL_RE = re.compile(r'^(\[)(.+)(\])$|^(\()(.+)(\))$')
+_LABEL_DIGIT_PREFIX_RE = re.compile(r'^(\d+)(.+)$')
 
 # How far before the first segment start to search for the bracket label inner content.
 # Needed because the parent SW match may start at "]" (skipping the preceding "[" and
@@ -467,6 +472,31 @@ def _try_bracket_label_match(  # pylint: disable=too-many-locals
     return new_start, new_end, new_blocks
 
 
+def _try_numeric_prefix_label_match(
+    token_index: _TokenIndex,
+    label_needle: str,
+    segments: List[Tuple[int, int]],
+) -> Optional[_MatchResult]:
+    """Match labels like "1-" when SW fails because the PDF has "1 -" (space between
+    digit and suffix).  Finds the digit prefix as an exact token via _exact_number_match
+    then extends the match to cover the immediately adjacent suffix characters."""
+    m = _LABEL_DIGIT_PREFIX_RE.match(label_needle)
+    if not m:
+        return None
+    numeric_part, suffix = m.group(1), m.group(2)
+    result = _exact_number_match(token_index, numeric_part, segments)
+    if result is None:
+        return None
+    num_start, num_end, num_blocks = result
+    haystack = token_index.haystack
+    pos = num_end
+    while pos < len(haystack) and haystack[pos] == ' ':
+        pos += 1
+    if haystack[pos:pos + len(suffix)] == suffix:
+        return num_start, pos + len(suffix), num_blocks + [(pos, pos + len(suffix))]
+    return result
+
+
 def _is_exact_sw_match(result: _MatchResult, needle_len: int) -> bool:
     """True when SW found the needle as one contiguous block (no gaps)."""
     _, _, blocks = result
@@ -567,6 +597,11 @@ def _fuzzy_match_field_value(  # pylint: disable=too-many-locals
         if bracket_match is not None:
             return bracket_match
 
+    if gap_match is None and field_value.sub_field_name == JatsSubFieldNames.REFERENCE_LABEL:
+        prefix_match = _try_numeric_prefix_label_match(token_index, needle, segments)
+        if prefix_match is not None:
+            return prefix_match
+
     return gap_match
 
 
@@ -587,7 +622,15 @@ def _search_range(
             fv.sub_field_name == JatsSubFieldNames.REFERENCE_LABEL
             and _is_pure_number(fv.text)
         )
-        pre = _SUB_FIELD_PARENT_PRE_BUFFER if is_pure_number_label else 0
+        if is_pure_number_label:
+            pre = _SUB_FIELD_PARENT_PRE_BUFFER
+        elif (
+            fv.sub_field_name == JatsSubFieldNames.REFERENCE_LABEL
+            and bool(_LABEL_DIGIT_PREFIX_RE.match(fv.text))
+        ):
+            pre = _SUB_FIELD_LABEL_DIGIT_PRE_BUFFER
+        else:
+            pre = 0
         return max(0, p_start - pre), 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
diff --git a/tests/training/jats/test_aligner.py b/tests/training/jats/test_aligner.py
index aa7600ca..3d437aaa 100644
--- a/tests/training/jats/test_aligner.py
+++ b/tests/training/jats/test_aligner.py
@@ -1045,3 +1045,39 @@ def test_pure_number_label_found_with_extra_text_before_parent(self):
         assert sub == JatsSubFieldNames.REFERENCE_LABEL, (
             f'"17" should be REFERENCE_LABEL even when parent starts 14 chars later; got {sub!r}'
         )
+
+    def test_single_digit_suffixed_label_found(self):
+        # Labels like "1-" fail SW matching because the PDF tokeniser produces "1 - …"
+        # (space between digit and dash), giving quality=0.5 < threshold.  The parent SW
+        # match also skips "1", so p_start lands at "-" (position 2).  The small
+        # digit-prefix pre-buffer and _try_numeric_prefix_label_match together fix this.
+        doc = _make_doc(
+            '1- Adebisi YA Sex workers should not be forgotten Am J Trop Med Hyg 2020',
+            '2- Li Q Early transmission dynamics N Engl J Med 2020',
+        )
+        fvs = [
+            _fv(
+                '1- Adebisi YA Sex workers should not be forgotten Am J Trop Med Hyg 2020',
+                JatsFieldNames.REFERENCE,
+            ),
+            _fv('1-', JatsFieldNames.REFERENCE, JatsSubFieldNames.REFERENCE_LABEL),
+            _fv('Adebisi YA', JatsFieldNames.REFERENCE, JatsSubFieldNames.REFERENCE_AUTHOR),
+            _fv('2020', JatsFieldNames.REFERENCE, JatsSubFieldNames.REFERENCE_YEAR),
+            _fv(
+                '2- Li Q Early transmission dynamics N Engl J Med 2020',
+                JatsFieldNames.REFERENCE,
+            ),
+            _fv('2-', JatsFieldNames.REFERENCE, JatsSubFieldNames.REFERENCE_LABEL),
+            _fv('Li Q', JatsFieldNames.REFERENCE, JatsSubFieldNames.REFERENCE_AUTHOR),
+            _fv('2020', JatsFieldNames.REFERENCE, JatsSubFieldNames.REFERENCE_YEAR),
+        ]
+        annotated = self._align(doc, fvs)
+        tokens = list(doc.iter_all_tokens())
+        ref1_label = next(t for t in tokens if t.text == '1')
+        assert annotated.get_token_sub_field(ref1_label) == JatsSubFieldNames.REFERENCE_LABEL, (
+            '"1" (label of first ref) should be REFERENCE_LABEL'
+        )
+        ref2_label = next(t for t in tokens if t.text == '2')
+        assert annotated.get_token_sub_field(ref2_label) == JatsSubFieldNames.REFERENCE_LABEL, (
+            '"2" (label of second ref) should be REFERENCE_LABEL'
+        )

From 08e97ffacfdd2be311bcaa490b1f7e63c764c29b Mon Sep 17 00:00:00 2001
From: Daniel Ecer <d.ecer@elifesciences.org>
Date: Tue, 30 Jun 2026 13:36:45 +0100
Subject: [PATCH 10/12] Fix first author token missing when consecutive
 reference SW matches overlap

  When _extend_match_for_needle_tail greedily extends a reference's parent
  match into the first token of the next reference, prev_parent_end can
  exceed the current reference's parent start (a_start). Without a guard,
  effective_prev_end is set past the token boundary, causing
  _label_tokens_for_blocks to skip it as a mid-token start.

  Apply the same <= a_start two-column guard to prev_parent_end that already
  existed for prev_id_end.
---
 sciencebeam_parser/training/jats/aligner.py | 119 ++++++++---
 tests/training/jats/test_aligner.py         | 206 ++++++++++++++++++++
 2 files changed, 300 insertions(+), 25 deletions(-)

diff --git a/sciencebeam_parser/training/jats/aligner.py b/sciencebeam_parser/training/jats/aligner.py
index 2ef05b76..63301730 100644
--- a/sciencebeam_parser/training/jats/aligner.py
+++ b/sciencebeam_parser/training/jats/aligner.py
@@ -94,6 +94,19 @@
 # aligns the parent by skipping the label digit, leaving p_start at the suffix char.
 # A small buffer is enough — single-digit + space = 2 chars.
 _SUB_FIELD_LABEL_DIGIT_PRE_BUFFER = 3
+# Author names appear before the title/journal anchor text that SW latches onto.
+# When the JATS given-names are initials ("RC") but the PDF has "R. C." (each initial
+# as a separate token), the run of gaps makes SW skip the author prefix entirely and
+# start the parent match at the title.  200 chars covers long multi-author lists too.
+_SUB_FIELD_REFERENCE_AUTHOR_PRE_BUFFER = 200
+# Identifier sub-fields (DOI, PMID, PMCID) are the only ones that appear in a URL tail
+# AFTER the parent match text.  Their match ends advance the backward-search floor so
+# the next reference's author search cannot reach back into the URL.
+_REFERENCE_IDENTIFIER_SUB_FIELDS = frozenset({
+    JatsSubFieldNames.REFERENCE_DOI,
+    JatsSubFieldNames.REFERENCE_PMID,
+    JatsSubFieldNames.REFERENCE_PMCID,
+})
 
 # Anchor+chain labelling strategy:
 # Smith-Waterman produces many tiny (1–4 char) matching blocks while traversing
@@ -265,6 +278,34 @@ def _match_quality(
     return matched / needle_len
 
 
+def _scan_tail_chars_at_token_starts(
+    tail: str,
+    scan_start: int,
+    token_index: '_TokenIndex',
+) -> int:
+    """Scan for each char in tail at token-start positions with gaps ≤ _MAX_HAYSTACK_GAP_TO_FILL.
+
+    Returns the absolute position just after the last successfully matched char.
+    If no chars can be matched, returns scan_start.  Used to bridge gaps such as
+    ". " between dotted initials ("L. Y.") when the needle has "ly".
+    """
+    pos = scan_start
+    end = scan_start
+    for char in tail:
+        limit = pos + _MAX_HAYSTACK_GAP_TO_FILL + 1
+        found = False
+        while pos < limit and pos < len(token_index.haystack):
+            if token_index.is_token_start(pos) and token_index.haystack[pos] == char:
+                end = pos + 1
+                pos = end
+                found = True
+                break
+            pos += 1
+        if not found:
+            break
+    return end
+
+
 def _extend_match_for_needle_tail(
     window: str,
     needle: str,
@@ -272,6 +313,7 @@ def _extend_match_for_needle_tail(
     abs_a_end: int,
     matched_blocks: List[Tuple[int, int, int]],
     abs_block_ranges: List[Tuple[int, int]],
+    token_index: Optional['_TokenIndex'] = None,
 ) -> Tuple[int, List[Tuple[int, int]]]:
     """Greedily extend the SW match to cover any unmatched needle suffix.
 
@@ -285,9 +327,12 @@ def _extend_match_for_needle_tail(
 
     Example: needle "brockmann d", match ends at "brockmann "; " , d" in the
     haystack has "d" at gap 2, which is within the fill threshold.
+
+    When token_index is provided, a second pass uses _scan_tail_chars_at_token_starts
+    to bridge larger gaps between remaining suffix chars (e.g. "l . y ." where
+    the needle suffix "ly" needs to hop over ". " to reach "y").
     """
-    last_needle_end = max(bi + size for _, bi, size in matched_blocks)
-    needle_tail = needle[last_needle_end:]
+    needle_tail = needle[max(bi + size for _, bi, size in matched_blocks):]
     if not needle_tail:
         return abs_a_end, abs_block_ranges
 
@@ -307,6 +352,15 @@ def _extend_match_for_needle_tail(
         return abs_a_end, abs_block_ranges
 
     ext_start = window_start + tail_start
+    # If chars remain unmatched and a token_index is available, try to bridge gaps
+    # to find them at token boundaries (handles ". " gaps between dotted initials).
+    if token_index is not None and match_count < len(needle_tail):
+        bridge_end = _scan_tail_chars_at_token_starts(
+            needle_tail[match_count:], ext_start + match_count, token_index
+        )
+        if bridge_end > ext_start + match_count:
+            return bridge_end, abs_block_ranges + [(ext_start, bridge_end)]
+
     return ext_start + match_count, abs_block_ranges + [(ext_start, ext_start + match_count)]
 
 
@@ -316,6 +370,7 @@ def _fuzzy_search_in_window(
     window_start: int,
     window_end: int,
     threshold: float,
+    token_index: Optional['_TokenIndex'] = None,
 ) -> Optional[_MatchResult]:
     """Try to find `needle` in haystack[window_start:window_end].
 
@@ -342,7 +397,7 @@ def _fuzzy_search_in_window(
         for ai, _bi, size in matched_blocks
     ]
     a_end, abs_block_ranges = _extend_match_for_needle_tail(
-        window, needle, window_start, a_end, matched_blocks, abs_block_ranges
+        window, needle, window_start, a_end, matched_blocks, abs_block_ranges, token_index
     )
     return a_start, a_end, abs_block_ranges
 
@@ -582,7 +637,9 @@ def _fuzzy_match_field_value(  # pylint: disable=too-many-locals
         start = seg_start
         while start < seg_end:
             end = min(start + window_size, seg_end)
-            result = _fuzzy_search_in_window(haystack, needle, start, end, config.threshold)
+            result = _fuzzy_search_in_window(
+                haystack, needle, start, end, config.threshold, token_index
+            )
             if result is not None:
                 if _is_exact_sw_match(result, need_len):
                     return result
@@ -613,25 +670,27 @@ def _search_range(
     front_matter_end: int,
     keywords_floor: int,
     reference_floor: int,
-    parent_match_by_field: Dict[str, Tuple[int, int]],
+    parent_match_by_field: Dict[str, Tuple[int, 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]
-        is_pure_number_label = (
-            fv.sub_field_name == JatsSubFieldNames.REFERENCE_LABEL
-            and _is_pure_number(fv.text)
-        )
-        if is_pure_number_label:
+        p_start, p_end, pre_parent_ref_floor = parent_match_by_field[fv.field_name]
+        if fv.sub_field_name == JatsSubFieldNames.REFERENCE_LABEL and _is_pure_number(fv.text):
             pre = _SUB_FIELD_PARENT_PRE_BUFFER
         elif (
             fv.sub_field_name == JatsSubFieldNames.REFERENCE_LABEL
             and bool(_LABEL_DIGIT_PREFIX_RE.match(fv.text))
         ):
             pre = _SUB_FIELD_LABEL_DIGIT_PRE_BUFFER
+        elif fv.sub_field_name == JatsSubFieldNames.REFERENCE_AUTHOR:
+            pre = _SUB_FIELD_REFERENCE_AUTHOR_PRE_BUFFER
         else:
             pre = 0
-        return max(0, p_start - pre), p_end + _SUB_FIELD_PARENT_BUFFER
+        # For authors: never extend before the end of the previous reference.
+        sub_start = max(p_start - pre, pre_parent_ref_floor) \
+            if fv.sub_field_name == JatsSubFieldNames.REFERENCE_AUTHOR \
+            else p_start - pre
+        return max(0, sub_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:
@@ -639,8 +698,7 @@ def _search_range(
         # 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
+        return (max(0, reference_floor - 200) if reference_floor > 0 else max(0, body_floor)), 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
@@ -728,16 +786,9 @@ def _extend_match_with_given_names_tail(
     abs_tail_start = fb_end + idx
     if not token_index.is_token_start(abs_tail_start):
         return match_range
-    match_count = 0
-    for i, char in enumerate(given_tail):
-        if abs_tail_start + i >= len(token_index.haystack):
-            break
-        if token_index.haystack[abs_tail_start + i] != char:
-            break
-        match_count += 1
-    if not match_count:
+    tail_end = _scan_tail_chars_at_token_starts(given_tail, abs_tail_start, token_index)
+    if tail_end == abs_tail_start:
         return match_range
-    tail_end = abs_tail_start + match_count
     return (
         match_range[0],
         max(match_range[1], tail_end),
@@ -968,13 +1019,16 @@ def align(  # pylint: disable=too-many-locals,too-many-branches,too-many-stateme
         front_matter_end = 0
         keywords_floor = 0
         reference_floor = 0
-        parent_match_by_field: Dict[str, Tuple[int, int]] = {}
+        parent_match_by_field: Dict[str, Tuple[int, int, int]] = {}
         missed_by_field: Dict[str, int] = {}
         matched_count = 0
         instance_by_field: Dict[str, int] = {}
         # Per-parent masked ranges: reset each time a new main-field match is
         # established so that sub-fields of one parent don't bleed into the next.
         sub_field_masked_ranges: Dict[str, List[Tuple[int, int]]] = {}
+        # Furthest end of any DOI/PMID/PMCID match for the current reference instance.
+        # Used to advance the backward-search floor past identifier URLs in the tail.
+        ref_id_subfield_end: Dict[str, int] = {}
 
         for fv in field_values:
             search_start, search_end = _search_range(
@@ -1108,7 +1162,18 @@ def align(  # pylint: disable=too-many-locals,too-many-branches,too-many-stateme
             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)
+                prev_parent_end = parent_match_by_field.get(fv.field_name, (0, 0, 0))[1]
+                prev_id_end = ref_id_subfield_end.pop(fv.field_name, 0)
+                # Advance the floor past identifier URLs (DOI/PMID/PMCID) that appear
+                # in the tail of the previous reference, beyond its parent match.
+                # Two-column guard: ignore either value if it falls at or after the
+                # current reference's parent start (handles overlapping SW matches and
+                # two-column layouts where the previous URL wraps past the current ref).
+                effective_prev_end = max(
+                    prev_parent_end if prev_parent_end <= a_start else 0,
+                    prev_id_end if prev_id_end <= a_start else 0,
+                )
+                parent_match_by_field[fv.field_name] = (a_start, a_end, effective_prev_end)
                 sub_field_masked_ranges[fv.field_name] = []
                 instance_by_field[fv.field_name] = (
                     instance_by_field.get(fv.field_name, 0) + 1
@@ -1117,6 +1182,10 @@ def align(  # pylint: disable=too-many-locals,too-many-branches,too-many-stateme
                 sub_field_masked_ranges.setdefault(fv.field_name, []).append(
                     (a_start, a_end)
                 )
+                if fv.sub_field_name in _REFERENCE_IDENTIFIER_SUB_FIELDS:
+                    ref_id_subfield_end[fv.field_name] = max(
+                        ref_id_subfield_end.get(fv.field_name, 0), a_end
+                    )
             instance_id = instance_by_field.get(fv.field_name, 0)
             _label_tokens_for_blocks(
                 annotated, token_index, block_ranges,
diff --git a/tests/training/jats/test_aligner.py b/tests/training/jats/test_aligner.py
index 3d437aaa..b13caeec 100644
--- a/tests/training/jats/test_aligner.py
+++ b/tests/training/jats/test_aligner.py
@@ -1081,3 +1081,209 @@ def test_single_digit_suffixed_label_found(self):
         assert annotated.get_token_sub_field(ref2_label) == JatsSubFieldNames.REFERENCE_LABEL, (
             '"2" (label of second ref) should be REFERENCE_LABEL'
         )
+
+    def test_reference_author_with_dotted_initials_in_pdf_annotated(self):
+        # JATS: "Smith AB" (initials without dots); PDF: "Smith, A. B." (dots after each initial)
+        doc = _make_doc(
+            'Smith, A. B. (2020). Some article title. Some Journal, 5, 7-9.',
+        )
+        fvs = [
+            _fv(
+                'Smith AB Some article title Some Journal 2020 5 7 9',
+                JatsFieldNames.REFERENCE,
+            ),
+            _fv('Smith', JatsFieldNames.REFERENCE, JatsSubFieldNames.REFERENCE_AUTHOR),
+        ]
+        annotated = self._align(doc, fvs)
+        tokens = list(doc.iter_all_tokens())
+        smith_token = next(t for t in tokens if t.text == 'Smith')
+        assert annotated.get_token_sub_field(smith_token) == JatsSubFieldNames.REFERENCE_AUTHOR
+
+    def test_reference_author_initials_annotated_when_pdf_has_compacted_token(self):
+        # JATS has "Jones AB" (two-letter initials run together); the PDF also renders
+        # them as a single compacted token "AB" rather than dotted "A. B.".  The initials
+        # token must be labeled even in compacted form — it must not be dropped because a
+        # mid-token SW block triggered a surname-only fallback.
+        doc = _make_doc(
+            'Smith A, Jones AB. Some article title. Some Journal 2020.',
+        )
+        fvs = [
+            _fv(
+                'Smith A Jones AB Some article title Some Journal 2020',
+                JatsFieldNames.REFERENCE,
+            ),
+            _fv('Smith A', JatsFieldNames.REFERENCE, JatsSubFieldNames.REFERENCE_AUTHOR),
+            _fv('Jones AB', JatsFieldNames.REFERENCE, JatsSubFieldNames.REFERENCE_AUTHOR),
+        ]
+        annotated = self._align(doc, fvs)
+        tokens = list(doc.iter_all_tokens())
+        ab_token = next(t for t in tokens if t.text == 'AB')
+        assert annotated.get_token_sub_field(ab_token) == JatsSubFieldNames.REFERENCE_AUTHOR, (
+            '"AB" is a compacted initials token; it must be labeled as REFERENCE_AUTHOR'
+        )
+
+    def test_reference_author_second_initial_labeled_when_primary_sw_drops_break_even_tail(self):
+        # "Jones-Smith AB" is the only author; the PDF renders initials as dotted separate
+        # tokens "Jones-Smith, A. B.".  SW matches the long surname (meeting the quality
+        # threshold without the initials) and drops "B" because adding it would decrease
+        # the alignment score.  The tail-bridging in _extend_match_for_needle_tail must
+        # bridge the ". " gap to find "B" at its token-start position.
+        doc = _make_doc(
+            'Jones-Smith, A. B. Some article title. Some Journal 2020.',
+        )
+        fvs = [
+            _fv(
+                'Jones-Smith AB Some article title Some Journal 2020',
+                JatsFieldNames.REFERENCE,
+            ),
+            JatsFieldValue(
+                text='Jones-Smith AB',
+                field_name=JatsFieldNames.REFERENCE,
+                sub_field_name=JatsSubFieldNames.REFERENCE_AUTHOR,
+                fallback_text='Jones-Smith',
+            ),
+        ]
+        annotated = self._align(doc, fvs)
+        tokens = list(doc.iter_all_tokens())
+        b_token = next(t for t in tokens if t.text == 'B')
+        assert annotated.get_token_sub_field(b_token) == JatsSubFieldNames.REFERENCE_AUTHOR, (
+            '"B" (second initial, after a ". " gap) must be labeled as REFERENCE_AUTHOR '
+            'when SW tail bridging bridges the ". " gap without a fallback match'
+        )
+
+    def test_reference_author_given_name_initials_labeled_when_surname_fallback_used(self):
+        # JATS has "Smíth AB" (accented í not in the PDF) which causes the primary match
+        # quality to fall below the threshold, triggering the surname-only fallback.  The
+        # PDF renders the initials as dotted separate tokens "A. B.".  The fallback tail
+        # extension must bridge the ". " gap between the two initials so both "A" and "B"
+        # are labeled as REFERENCE_AUTHOR — not just the first.
+        doc = _make_doc(
+            'Smith, A. B. Some article title. Some Journal 2020.',
+        )
+        fvs = [
+            _fv(
+                'Smíth AB Some article title Some Journal 2020',
+                JatsFieldNames.REFERENCE,
+            ),
+            JatsFieldValue(
+                text='Smíth AB',
+                field_name=JatsFieldNames.REFERENCE,
+                sub_field_name=JatsSubFieldNames.REFERENCE_AUTHOR,
+                fallback_text='Smíth',
+            ),
+        ]
+        annotated = self._align(doc, fvs)
+        tokens = list(doc.iter_all_tokens())
+        b_token = next(t for t in tokens if t.text == 'B')
+        assert annotated.get_token_sub_field(b_token) == JatsSubFieldNames.REFERENCE_AUTHOR, (
+            '"B" (second initial, after a ". " gap) must be labeled as REFERENCE_AUTHOR '
+            'even when the surname-only fallback is used'
+        )
+
+    def test_prior_reference_doi_subfield_end_advances_backward_search_floor(self):
+        # The next reference's PDF text ("Smith XY 2002") contains no "Doe", so the
+        # primary match fails and the 3-char fallback "Doe" is tried.  Without the
+        # floor fix, "doe" in the prior reference's URL would be matched.
+        doc = _make_doc(
+            'Jones AB 2001. https://doe.org/10.1/abc',
+            'Smith XY 2002.',
+        )
+        fvs = [
+            _fv('Jones AB 2001', JatsFieldNames.REFERENCE),
+            _fv('Jones AB', JatsFieldNames.REFERENCE, JatsSubFieldNames.REFERENCE_AUTHOR),
+            _fv('10.1/abc', JatsFieldNames.REFERENCE, JatsSubFieldNames.REFERENCE_DOI),
+            _fv('Smith XY 2002', JatsFieldNames.REFERENCE),
+            JatsFieldValue(
+                text='Doe XY',
+                field_name=JatsFieldNames.REFERENCE,
+                sub_field_name=JatsSubFieldNames.REFERENCE_AUTHOR,
+                fallback_text='Doe',
+            ),
+        ]
+        annotated = self._align(doc, fvs)
+        tokens = list(doc.iter_all_tokens())
+        doe_url_token = next(t for t in tokens if t.text == 'doe')
+        assert annotated.get_token_sub_field(doe_url_token) != JatsSubFieldNames.REFERENCE_AUTHOR, (
+            '"doe" in the prior reference URL must not be labeled as REFERENCE_AUTHOR; '
+            'the DOI sub-field end should advance the floor past the URL'
+        )
+
+    def test_reference_author_labeled_when_prev_ref_sw_match_overlaps_current_ref_start(self):
+        # Ref 1's JATS text ends with "Zorba" — the same word that starts ref 2.
+        # The SW match for ref 1's parent consumes "Zorba" from the haystack, making
+        # prev_parent_end land PAST ref 2's parent start.  Without the guard, the
+        # backward-search floor is set beyond "Zorba" so the token is never labeled.
+        doc = _make_doc(
+            'Vreeken AB 2001 Source A',
+            'Zorba BC 2002 Title',
+        )
+        fvs = [
+            _fv('Vreeken AB 2001 Source A Zorba', JatsFieldNames.REFERENCE),
+            _fv('Vreeken AB', JatsFieldNames.REFERENCE, JatsSubFieldNames.REFERENCE_AUTHOR),
+            _fv('Zorba BC 2002 Title', JatsFieldNames.REFERENCE),
+            JatsFieldValue(
+                text='Zorba BC',
+                field_name=JatsFieldNames.REFERENCE,
+                sub_field_name=JatsSubFieldNames.REFERENCE_AUTHOR,
+                fallback_text='Zorba',
+            ),
+        ]
+        annotated = self._align(doc, fvs)
+        tokens = list(doc.iter_all_tokens())
+        zorba_token = next(t for t in tokens if t.text == 'Zorba')
+        assert annotated.get_token_sub_field(zorba_token) == JatsSubFieldNames.REFERENCE_AUTHOR, (
+            '"Zorba" is ref 2\'s author; the prev_parent_end guard must prevent the '
+            'floor from being placed past the current reference\'s start position'
+        )
+
+    def test_prior_reference_et_al_not_labeled_by_next_reference_author_search(self):
+        # Reference 1's PDF has "et al." but reference 1's JATS does not list it as
+        # an author, so it should remain unlabeled.  Reference 2 follows with dotted
+        # initials in the PDF ("Smith, B. C.") and "et al." in its JATS.  Reference
+        # 2's author search must not reach back and label reference 1's "et al." tokens.
+        doc = _make_doc(
+            'Jones, A. et al. Title A Source A 2001',
+            'Smith, B. C. Title B Source B 2002',
+        )
+        fvs = [
+            _fv('Jones A et al. Title A Source A 2001', JatsFieldNames.REFERENCE),
+            _fv('Jones A', JatsFieldNames.REFERENCE, JatsSubFieldNames.REFERENCE_AUTHOR),
+            _fv('Smith BC Title B Source B 2002', JatsFieldNames.REFERENCE),
+            _fv('Smith BC', JatsFieldNames.REFERENCE, JatsSubFieldNames.REFERENCE_AUTHOR),
+            _fv('et al.', JatsFieldNames.REFERENCE, JatsSubFieldNames.REFERENCE_AUTHOR),
+        ]
+        annotated = self._align(doc, fvs)
+        tokens = list(doc.iter_all_tokens())
+        et_token = next(t for t in tokens if t.text == 'et')
+        al_token = next(t for t in tokens if t.text == 'al')
+        assert annotated.get_token_sub_field(et_token) != JatsSubFieldNames.REFERENCE_AUTHOR, (
+            '"et al." belongs to reference 1 whose JATS has no "et al." author — '
+            'it must not be labeled by reference 2\'s author search'
+        )
+        assert annotated.get_token_sub_field(al_token) != JatsSubFieldNames.REFERENCE_AUTHOR
+
+    def test_reference_author_found_when_prev_ref_source_appears_later_in_haystack(self):
+        # Simulates a two-column PDF layout: ref 1's source text physically appears in the
+        # token stream AFTER ref 2's author tokens.  Before this fix, reference_floor
+        # advanced past ref 2's author position when ref 1's source was matched, causing
+        # the author search to miss ref 2's author.  The fix uses the previous reference's
+        # parent-match end as the backward floor instead of the accumulated reference_floor.
+        doc = _make_doc(
+            'Jones A 2001',       # ref 1 parent (short, appears first in reading order)
+            'Smith BC 2002',      # ref 2 parent + author (follows ref 1 in reading order)
+            'source of ref 1',    # ref 1 source (appears last — two-column reading order)
+        )
+        fvs = [
+            _fv('Jones A 2001', JatsFieldNames.REFERENCE),
+            _fv('Jones A', JatsFieldNames.REFERENCE, JatsSubFieldNames.REFERENCE_AUTHOR),
+            _fv('source of ref 1', JatsFieldNames.REFERENCE, JatsSubFieldNames.REFERENCE_SOURCE),
+            _fv('Smith BC 2002', JatsFieldNames.REFERENCE),
+            _fv('Smith BC', JatsFieldNames.REFERENCE, JatsSubFieldNames.REFERENCE_AUTHOR),
+        ]
+        annotated = self._align(doc, fvs)
+        tokens = list(doc.iter_all_tokens())
+        smith_token = next(t for t in tokens if t.text == 'Smith')
+        assert annotated.get_token_sub_field(smith_token) == JatsSubFieldNames.REFERENCE_AUTHOR, (
+            '"Smith" is ref 2\'s first author; it must be annotated even though '
+            'ref 1\'s source match appears after it in the token stream'
+        )

From 8798fa28d849233703e3219b0b9e890c358896f9 Mon Sep 17 00:00:00 2001
From: Daniel Ecer <d.ecer@elifesciences.org>
Date: Tue, 30 Jun 2026 21:25:41 +0100
Subject: [PATCH 11/12] Attach trailing period when it falls outside the outer
 SW match range

  When the gap between the last author initial and the next SW match
  point exceeds the fill threshold, the abbreviation period immediately
  following that initial has no entry in token_label_by_id. The
  trailing-period pass previously required entry is not None, so it
  silently skipped the period. Relaxing the condition to also handle
  entry is None attaches it correctly.
---
 sciencebeam_parser/training/jats/aligner.py |  3 +--
 tests/training/jats/test_aligner.py         | 19 +++++++++++++++++++
 2 files changed, 20 insertions(+), 2 deletions(-)

diff --git a/sciencebeam_parser/training/jats/aligner.py b/sciencebeam_parser/training/jats/aligner.py
index 63301730..9df2b9ad 100644
--- a/sciencebeam_parser/training/jats/aligner.py
+++ b/sciencebeam_parser/training/jats/aligner.py
@@ -947,8 +947,7 @@ def _attach_sub_field_trailing_periods(
             last_instance = entry[2]
         elif (
             last_was_sub
-            and entry is not None
-            and entry[0] == field_name
+            and (entry is None or entry[0] == field_name)
             and normalize_for_alignment(token.text or '') == '.'
         ):
             annotated.set_token_label(token, field_name, sub_field_name, last_instance)
diff --git a/tests/training/jats/test_aligner.py b/tests/training/jats/test_aligner.py
index b13caeec..23ff5d17 100644
--- a/tests/training/jats/test_aligner.py
+++ b/tests/training/jats/test_aligner.py
@@ -415,6 +415,25 @@ def test_trailing_period_after_last_initial_is_labeled(self):
             f'{annotated.get_token_sub_field(period_token)!r}'
         )
 
+    def test_trailing_period_unlabeled_by_outer_match_is_attached(self):
+        # When the gap between the last initial and the next matched word exceeds
+        # _MAX_HAYSTACK_GAP_TO_FILL, the abbreviation period is outside the outer
+        # SW match range (entry=None). The trailing-period pass must still attach it.
+        doc = _make_doc('Zorblax, C. A. (2010). Some title.')
+        fvs = [
+            _fv('Zorblax CA Some title 2010', JatsFieldNames.REFERENCE),
+            _fv('Zorblax CA', JatsFieldNames.REFERENCE, JatsSubFieldNames.REFERENCE_AUTHOR),
+        ]
+        annotated = self._align(doc, fvs)
+        tokens = list(doc.iter_all_tokens())
+        a_idx = next(i for i, t in enumerate(tokens) if normalize_for_alignment(t.text) == 'a')
+        period_token = tokens[a_idx + 1]
+        assert period_token.text == '.', 'Expected period token after A'
+        assert annotated.get_token_sub_field(period_token) == JatsSubFieldNames.REFERENCE_AUTHOR, (
+            f'Period after A should be REFERENCE_AUTHOR, got '
+            f'{annotated.get_token_sub_field(period_token)!r}'
+        )
+
     def test_gap_fill_merges_per_name_author_spans(self):
         # Between two per-name author matches, separator tokens (comma, semicolon,
         # initials with periods) should be filled in as REFERENCE_AUTHOR so that

From 5e5a6ea89f7584ae58f0b83d29aa816db700fd4b Mon Sep 17 00:00:00 2001
From: Daniel Ecer <d.ecer@elifesciences.org>
Date: Wed, 1 Jul 2026 17:04:36 +0100
Subject: [PATCH 12/12] Extend backward pre-buffer to reference-source
 sub-field
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

  When a PDF orders a reference as author → source → article-title but the
  JATS element order is author → article-title → source, Smith-Waterman
  anchors the parent match on the article-title. The resulting p_start lands
  after the source in the haystack, so the source sub-field search (started
  at p_start) misses it entirely.

  The fix adds REFERENCE_SOURCE to the elif branch that sets the 200-char
  backward pre-buffer, already used by REFERENCE_AUTHOR. Both sub-fields
  share the same pre_parent_ref_floor guard that prevents the extended search
  from reaching into a preceding reference's matched range.
---
 sciencebeam_parser/training/jats/aligner.py | 17 +++++-
 tests/training/jats/test_aligner.py         | 66 +++++++++++++++++++++
 2 files changed, 80 insertions(+), 3 deletions(-)

diff --git a/sciencebeam_parser/training/jats/aligner.py b/sciencebeam_parser/training/jats/aligner.py
index 9df2b9ad..6e853e83 100644
--- a/sciencebeam_parser/training/jats/aligner.py
+++ b/sciencebeam_parser/training/jats/aligner.py
@@ -682,13 +682,24 @@ def _search_range(
             and bool(_LABEL_DIGIT_PREFIX_RE.match(fv.text))
         ):
             pre = _SUB_FIELD_LABEL_DIGIT_PRE_BUFFER
-        elif fv.sub_field_name == JatsSubFieldNames.REFERENCE_AUTHOR:
+        elif fv.sub_field_name in {
+            JatsSubFieldNames.REFERENCE_AUTHOR,
+            JatsSubFieldNames.REFERENCE_SOURCE,
+        }:
+            # Source can precede the parent's SW match start when the PDF orders
+            # source before article-title but the JATS parent text has them reversed.
+            # SW then latches onto the article-title anchor and sets p_start after
+            # the source, so we need the same backward buffer as for authors.
             pre = _SUB_FIELD_REFERENCE_AUTHOR_PRE_BUFFER
         else:
             pre = 0
-        # For authors: never extend before the end of the previous reference.
+        # For authors and source: never extend before the end of the previous
+        # reference (prevents sub-field matches from bleeding into earlier bibls).
         sub_start = max(p_start - pre, pre_parent_ref_floor) \
-            if fv.sub_field_name == JatsSubFieldNames.REFERENCE_AUTHOR \
+            if fv.sub_field_name in {
+                JatsSubFieldNames.REFERENCE_AUTHOR,
+                JatsSubFieldNames.REFERENCE_SOURCE,
+            } \
             else p_start - pre
         return max(0, sub_start), p_end + _SUB_FIELD_PARENT_BUFFER
     if fv.field_name in _BODY_CONTENT_FIELDS:
diff --git a/tests/training/jats/test_aligner.py b/tests/training/jats/test_aligner.py
index 23ff5d17..cfdd482a 100644
--- a/tests/training/jats/test_aligner.py
+++ b/tests/training/jats/test_aligner.py
@@ -1281,6 +1281,72 @@ def test_prior_reference_et_al_not_labeled_by_next_reference_author_search(self)
         )
         assert annotated.get_token_sub_field(al_token) != JatsSubFieldNames.REFERENCE_AUTHOR
 
+    def test_reference_source_found_when_it_precedes_parent_match_start(self):
+        # The JATS parent text for a reference concatenates sub-fields in JATS element
+        # order (author → article-title → source).  The PDF text may instead have the
+        # source immediately after the author and before the article-title.  SW then
+        # latches onto the article-title anchor and produces a parent match start that
+        # is AFTER the source in the haystack, causing the source sub-field search
+        # (anchored at p_start) to miss the source.
+        #
+        # Reproduces the GEOCAPES pattern: institutional author whose name also starts
+        # the source, appearing in the order author → source → article-title in the PDF.
+        doc = _make_doc(
+            # preceding reference (establishes pre_parent_ref_floor)
+            'Zorblax A 2001 Title of Prev Article Source of Prev',
+            # GEOCAPES-like reference: author → source → article-title in PDF text
+            'Zorbax. Zorbax Research Group Source. Article Title Here.',
+        )
+        fvs = [
+            # preceding reference
+            _fv('Zorblax A 2001 Title of Prev Article Source of Prev',
+                JatsFieldNames.REFERENCE),
+            _fv('Zorblax A', JatsFieldNames.REFERENCE, JatsSubFieldNames.REFERENCE_AUTHOR),
+            _fv('Source of Prev', JatsFieldNames.REFERENCE, JatsSubFieldNames.REFERENCE_SOURCE),
+            # GEOCAPES-like reference: JATS order is author → article-title → source
+            # but PDF text order is author → source → article-title
+            _fv('Zorbax Article Title Here',  # parent: author + article-title (no source)
+                JatsFieldNames.REFERENCE),
+            _fv('Zorbax', JatsFieldNames.REFERENCE, JatsSubFieldNames.REFERENCE_AUTHOR),
+            _fv('Article Title Here', JatsFieldNames.REFERENCE,
+                JatsSubFieldNames.REFERENCE_ARTICLE_TITLE),
+            _fv('Zorbax Research Group Source', JatsFieldNames.REFERENCE,
+                JatsSubFieldNames.REFERENCE_SOURCE),
+        ]
+        annotated = self._align(doc, fvs)
+        tokens = list(doc.iter_all_tokens())
+        source_token = next(t for t in tokens if t.text == 'Research')
+        assert annotated.get_token_sub_field(source_token) == JatsSubFieldNames.REFERENCE_SOURCE, (
+            '"Research" is part of the source that precedes the parent SW match start; '
+            'the backward pre-buffer must allow the source search to reach it'
+        )
+        author_token = next(t for t in tokens if t.text == 'Zorbax')
+        assert annotated.get_token_sub_field(author_token) == JatsSubFieldNames.REFERENCE_AUTHOR, (
+            '"Zorbax" is the author; REFERENCE_AUTHOR must retain its backward pre-buffer '
+            'so the author is not displaced by the source pre-buffer fix'
+        )
+
+    def test_reference_author_found_when_parent_match_starts_after_author(self):
+        # When the JATS parent text contains only the article-title (no author), the parent
+        # SW match anchors at the title, making p_start land after the author in the haystack.
+        # The REFERENCE_AUTHOR backward pre-buffer must extend the search back to reach it.
+        # This reproduces the GEOCAPES regression where adding a REFERENCE_SOURCE pre-buffer
+        # accidentally zeroed REFERENCE_AUTHOR's pre-buffer by removing its elif branch.
+        doc = _make_doc('Zorblax Very Long Article Title Here End.')
+        fvs = [
+            _fv('Very Long Article Title Here End', JatsFieldNames.REFERENCE),
+            _fv('Zorblax', JatsFieldNames.REFERENCE, JatsSubFieldNames.REFERENCE_AUTHOR),
+            _fv('Very Long Article Title Here End', JatsFieldNames.REFERENCE,
+                JatsSubFieldNames.REFERENCE_ARTICLE_TITLE),
+        ]
+        annotated = self._align(doc, fvs)
+        tokens = list(doc.iter_all_tokens())
+        author_token = next(t for t in tokens if t.text == 'Zorblax')
+        assert annotated.get_token_sub_field(author_token) == JatsSubFieldNames.REFERENCE_AUTHOR, (
+            'Author precedes the parent SW match start; the backward pre-buffer must allow '
+            'the author search to reach back to it'
+        )
+
     def test_reference_author_found_when_prev_ref_source_appears_later_in_haystack(self):
         # Simulates a two-column PDF layout: ref 1's source text physically appears in the
         # token stream AFTER ref 2's author tokens.  Before this fix, reference_floor