From 6c8d64e4dd1674444324a8c108a151dfd1a891f7 Mon Sep 17 00:00:00 2001 From: Daniel Ecer Date: Thu, 20 Aug 2026 10:04:24 +0100 Subject: [PATCH 1/4] Record the label cardinality of generated training data Nothing checked that a generated document's labels agreed with the JATS they were aligned from. The reference segmenter trained on a corpus where roughly half the documents presented their whole reference list as one reference, and the only sign of it was a coverage ratio in a log line; finding it meant counting labels in the artifact afterwards. Every run now writes a quality.jsonl per model it generated for, one row per source document, holding the count at each stage where the cardinality can change: the references the JATS declares, how many of them the aligner placed, and what the model wrote per entity. The reference of record is the set JatsFieldExtractor emits - back/ref-list/ref with non-empty text - since a peer-review sub-article's own ref-list is never offered to the aligner, and counting every ref-list instead reports a shortfall the pipeline never had: 1708 references for ore against the 1679 it is given. Per model rather than per corpus because generation is run per model. A corpus commonly holds one model's data at one document set and another model's at a different one - the committed corpora carry the reference segmenter and citation at 38 and 50 documents against 5 and 10 for segmentation - so one file per corpus would describe the last run rather than the data beside it. The citation model takes a per-label count rather than a third entity count, because its parser's root element path is bibl, so every element is its own training sequence and the start count equals the element count by construction. Presence per reference on both sides, which holds whichever convention a label uses: covers a whole author list where is written once per page number. The two sides differ legitimately - ore marks no identifier at all because 2% of its printed references carry a DOI-like string while its JATS carries one for 1205 of 1679 - so this is a rate to compare across regenerations rather than something to expect. The record is written by the parent process as each document finishes, so a document that timed out or died is present in every model's file with its status rather than missing, a document that produced no file at all is visible where nothing iterating generated output would show it, and an interrupted run keeps what it had. A JATS that will not parse is its own status, distinct from one that declares no references. Measuring only: no threshold, no verdict, and nothing is filtered. Over the committed corpora the counts come out at 1679 references against 1678 elements for ore and 1674 against 1648 for scielo_preprints-jats, naming both documents whose JATS declares no references, the one whose JATS will not parse, PPR459453 at 45 references to 2 elements, and the seven documents holding more elements than their JATS has references. --- doc/training.md | 52 +++ .../training/cli/generate_data.py | 308 ++++++++++++++---- .../training/jats/annotated_document.py | 13 + .../training/jats/field_extractor.py | 27 +- .../training/quality/__init__.py | 0 .../training/quality/counting.py | 110 +++++++ sciencebeam_parser/training/quality/record.py | 187 +++++++++++ tests/training/cli/generate_data_test.py | 174 ++++++++++ tests/training/jats/test_field_extractor.py | 68 +++- tests/training/quality/__init__.py | 0 tests/training/quality/counting_test.py | 122 +++++++ tests/training/quality/record_test.py | 184 +++++++++++ 12 files changed, 1179 insertions(+), 66 deletions(-) create mode 100644 sciencebeam_parser/training/quality/__init__.py create mode 100644 sciencebeam_parser/training/quality/counting.py create mode 100644 sciencebeam_parser/training/quality/record.py create mode 100644 tests/training/quality/__init__.py create mode 100644 tests/training/quality/counting_test.py create mode 100644 tests/training/quality/record_test.py diff --git a/doc/training.md b/doc/training.md index 671e3ddc..68510590 100644 --- a/doc/training.md +++ b/doc/training.md @@ -66,6 +66,58 @@ python -m sciencebeam_parser.training.cli.generate_data \ Additionally the `--gzip` argument can be passed in, resulting in gzip (`.gz`) compressed output files. +#### The quality record + +Every run writes a `quality.jsonl` per model it generated for, one JSON line per +source document, whether the document succeeded or not. It holds the count at each +stage where the cardinality of the labels can change, so that a corpus can be +compared against the JATS it was aligned from without counting labels in the +generated data afterwards. + +The record is per model because generation is run per model: a corpus commonly +holds one model's data at one document set and another model's at a different one, +and a record covering the whole corpus would describe the last run rather than the +data beside it. With `--use-directory-structure` each file sits beside that model's +`corpus` directory, otherwise it is `.quality.jsonl` in the output path: + +```text +reference-segmenter/quality.jsonl +citation/quality.jsonl +``` + +```json +{ + "document_id": "PPR459453", + "source_filename": "PPR459453.pdf", + "status": "ok", + "model": "citation", + "jats": {"status": "ok", "reference_count": 45, "aligned_reference_count": 2}, + "written": true, + "entity_element_count": 2, + "label_counts": {"": {"jats": 44, "marked": 2}} +} +``` + +- `jats.status` is `ok`, `missing` (no JATS was matched), `unparsable` or + `unreadable`. A `reference_count` of 0 with status `ok` is a JATS that declares + no references — there was never anything to align. +- `aligned_reference_count` is how many of those references the aligner placed, so + the difference from `reference_count` is alignment's. +- `entity_element_count` is what the model wrote per entity: `bibl` for + `reference-segmenter` and `citation`, and absent for a model whose labels mark + regions rather than repeated entities. `written: false` is a model that found no + entities and so wrote no file at all. +- `label_counts` is per citation label, over references rather than occurrences: + `jats` counts references whose JATS carries a sub-field for that label, `marked` + counts references the training data marks it in. The two differ legitimately — + a printed reference does not carry everything its JATS does, so a low rate for + an identifier or a URL is usually the page rather than the pipeline. + +The record is written by the parent process as each document finishes, so a run +that is interrupted keeps the records it had, and a document that timed out or +failed is present with a `status` of `timeout` or `error` in every model's file +rather than missing. + ### Annotating `tei` training data for the sequence models After the `tei` training data has been generated, it should get reviewed and manually annotated. diff --git a/sciencebeam_parser/training/cli/generate_data.py b/sciencebeam_parser/training/cli/generate_data.py index 4b46cf23..9d2e0236 100644 --- a/sciencebeam_parser/training/cli/generate_data.py +++ b/sciencebeam_parser/training/cli/generate_data.py @@ -8,7 +8,17 @@ import time from dataclasses import dataclass, field -from typing import Callable, Dict, Iterable, List, NamedTuple, Optional, Sequence, Tuple +from typing import ( + Callable, + Dict, + FrozenSet, + Iterable, + List, + NamedTuple, + Optional, + Sequence, + Tuple +) from lxml import etree @@ -57,7 +67,24 @@ JatsFieldNames, JatsSubFieldNames, ) -from sciencebeam_parser.training.jats.field_extractor import JatsFieldExtractor +from sciencebeam_parser.training.jats.field_extractor import ( + JatsFieldExtractor, + iter_reference_sub_field_names +) +from sciencebeam_parser.training.quality.counting import ( + ENTITY_ELEMENT_NAME_BY_MODEL, + count_citation_labels, + count_entity_elements +) +from sciencebeam_parser.training.quality.record import ( + DocumentQualityRecord, + DocumentStatus, + JatsQualityRecord, + JatsStatus, + ModelQualityRecord, + QualityRecordWriter, + get_failed_document_quality_record +) from sciencebeam_parser.training.jats.aligner import LayoutDocumentJatsAligner from sciencebeam_parser.training.jats.segmentation import SegmentationLabelDeriver @@ -245,6 +272,7 @@ class TrainingDataDocumentContext(NamedTuple): gzip_enabled: bool jats_annotated_document: Optional[JatsAnnotatedLayoutDocument] = None jats_segmentation_labels: Optional[Dict[int, str]] = None + jats_reference_sub_field_names: Sequence[FrozenSet[str]] = () @property def source_name(self) -> str: @@ -467,11 +495,19 @@ def get_default_tei_sub_directory( ) -> Optional[str]: return tei_training_data_generator.get_default_tei_sub_directory() + def get_quality_label_counts( # pylint: disable=unused-argument + self, + model_data_list_list: Sequence[Sequence[LayoutModelData]], + document_context: TrainingDataDocumentContext + ) -> Optional[Dict[str, Dict[str, int]]]: + """Per-label counts for a model whose entity cardinality cannot change.""" + return None + def generate_data_for_layout_document( self, layout_document: LayoutDocument, document_context: TrainingDataDocumentContext - ): + ) -> ModelQualityRecord: tei_training_data_generator = self.get_tei_training_data_generator(document_context) tei_file_path = self._get_file_path_with_suffix( tei_training_data_generator.get_default_tei_filename_suffix(), @@ -490,7 +526,13 @@ def generate_data_for_layout_document( )) if not model_data_list_list: LOGGER.info('no entities found, skipping (%r)', tei_file_path) - return + return ModelQualityRecord( + model_name=self.model_name, + written=False, + entity_element_count=( + 0 if self.model_name in ENTITY_ELEMENT_NAME_BY_MODEL else None + ), + ) training_tei_root = ( tei_training_data_generator .get_training_tei_xml_for_multiple_model_data_iterables( @@ -511,6 +553,14 @@ def generate_data_for_layout_document( ), encoding='utf-8' ) + return ModelQualityRecord( + model_name=self.model_name, + written=True, + entity_element_count=count_entity_elements(self.model_name, training_tei_root), + label_counts=self.get_quality_label_counts( + model_data_list_list, document_context + ), + ) class AbstractDocumentModelTrainingDataGenerator(AbstractModelTrainingDataGenerator): @@ -1166,6 +1216,20 @@ class CitationModelTrainingDataGenerator(AbstractDocumentModelTrainingDataGenera def get_main_model(self, document_context: TrainingDataDocumentContext) -> Model: return document_context.fulltext_models.citation_model + def get_quality_label_counts( + self, + model_data_list_list: Sequence[Sequence[LayoutModelData]], + document_context: TrainingDataDocumentContext + ) -> Optional[Dict[str, Dict[str, int]]]: + # Every <bibl> is its own training sequence, so this model's entity count + # cannot change again after the TEI; which labels are marked can. + if not document_context.jats_reference_sub_field_names: + return None + return count_citation_labels( + document_context.jats_reference_sub_field_names, + model_data_list_list + ) + def get_jats_label_fn(self) -> Optional[JatsLabelFn]: previous_identifier: Optional[Tuple[int, Optional[str]]] = None @@ -1266,19 +1330,67 @@ def _select_generators( return selected +def get_enabled_model_names(enabled_models: Optional[frozenset]) -> List[str]: + return [ + training_data_generator.model_name + for training_data_generator in _select_generators(enabled_models) + ] + + +class JatsAnnotationResult(NamedTuple): + status: str + annotated_document: Optional[JatsAnnotatedLayoutDocument] = None + reference_sub_field_names: Sequence[FrozenSet[str]] = () + + @property + def reference_count(self) -> Optional[int]: + if self.status != JatsStatus.OK: + return None + return len(self.reference_sub_field_names) + + def _build_jats_annotations( layout_document: LayoutDocument, jats_xml_filename: str, -) -> Optional[JatsAnnotatedLayoutDocument]: +) -> JatsAnnotationResult: try: with auto_download_input_file(jats_xml_filename, auto_decompress=True) as local_xml: root = etree.parse(local_xml).getroot() + except etree.XMLSyntaxError: + LOGGER.warning('JATS XML could not be parsed: %r', jats_xml_filename, exc_info=True) + return JatsAnnotationResult(status=JatsStatus.UNPARSABLE) except Exception: # pylint: disable=broad-except LOGGER.warning('Failed to load JATS XML: %r', jats_xml_filename, exc_info=True) - return None + return JatsAnnotationResult(status=JatsStatus.UNREADABLE) field_values = list(JatsFieldExtractor().iter_field_values(root)) LOGGER.debug('JATS field values count: %d', len(field_values)) - return LayoutDocumentJatsAligner().align(layout_document, field_values) + return JatsAnnotationResult( + status=JatsStatus.OK, + annotated_document=LayoutDocumentJatsAligner().align(layout_document, field_values), + reference_sub_field_names=list(iter_reference_sub_field_names(root)), + ) + + +def _get_document_quality_record( + document_context: TrainingDataDocumentContext, + jats_result: JatsAnnotationResult, + model_records: Sequence[ModelQualityRecord], +) -> DocumentQualityRecord: + annotated_document = jats_result.annotated_document + return DocumentQualityRecord( + document_id=document_context.source_name, + source_filename=document_context.source_filename, + jats=JatsQualityRecord( + status=jats_result.status, + reference_count=jats_result.reference_count, + aligned_reference_count=( + annotated_document.get_aligned_instance_count(JatsFieldNames.REFERENCE) + if annotated_document is not None + else None + ), + ), + models=model_records, + ) def generate_training_data_for_layout_document( @@ -1293,18 +1405,18 @@ def generate_training_data_for_layout_document( gzip_enabled: bool = False, jats_xml_filename: Optional[str] = None, enabled_models: Optional[frozenset] = None, -): - model_result_cache = ModelResultCache() - jats_annotated: Optional[JatsAnnotatedLayoutDocument] = None +) -> DocumentQualityRecord: + jats_result = JatsAnnotationResult(status=JatsStatus.MISSING) jats_seg_labels: Optional[Dict[int, str]] = None if jats_xml_filename: - jats_annotated = _build_jats_annotations(layout_document, jats_xml_filename) - if jats_annotated: + jats_result = _build_jats_annotations(layout_document, jats_xml_filename) + if jats_result.annotated_document: jats_seg_labels = SegmentationLabelDeriver().derive_labels( - layout_document, jats_annotated + layout_document, jats_result.annotated_document ) LOGGER.debug( - 'JATS coverage ratio: %.2f', jats_annotated.coverage_ratio() + 'JATS coverage ratio: %.2f', + jats_result.annotated_document.coverage_ratio() ) document_context = TrainingDataDocumentContext( output_path=output_path, @@ -1313,16 +1425,23 @@ def generate_training_data_for_layout_document( fulltext_models=fulltext_models, use_model=use_model, use_directory_structure=use_directory_structure, - model_result_cache=model_result_cache, + model_result_cache=ModelResultCache(), gzip_enabled=gzip_enabled, - jats_annotated_document=jats_annotated, + jats_annotated_document=jats_result.annotated_document, jats_segmentation_labels=jats_seg_labels, + jats_reference_sub_field_names=jats_result.reference_sub_field_names, + ) + return _get_document_quality_record( + document_context=document_context, + jats_result=jats_result, + model_records=[ + training_data_generator.generate_data_for_layout_document( + layout_document=layout_document, + document_context=document_context + ) + for training_data_generator in _select_generators(enabled_models) + ], ) - for training_data_generator in _select_generators(enabled_models): - training_data_generator.generate_data_for_layout_document( - layout_document=layout_document, - document_context=document_context - ) def get_layout_document_for_source_filename( @@ -1366,7 +1485,7 @@ def generate_training_data_for_source_filename( gzip_enabled: bool, xml_file_list: Optional[Sequence[str]] = None, enabled_models: Optional[frozenset] = None, -): +) -> DocumentQualityRecord: LOGGER.debug('use_model: %r', use_model) layout_document = get_layout_document_for_source_filename( source_filename, @@ -1379,7 +1498,7 @@ def generate_training_data_for_source_filename( LOGGER.info('Using JATS XML: %r', jats_xml_filename) else: LOGGER.warning('No matching JATS XML found for: %r', source_filename) - generate_training_data_for_layout_document( + return generate_training_data_for_layout_document( layout_document=layout_document, output_path=output_path, source_filename=source_filename, @@ -1451,10 +1570,11 @@ def _worker_init() -> None: _worker_sciencebeam_parser = ScienceBeamParser.from_config(config) -def _worker_process(kwargs: dict) -> bool: +def _worker_process(kwargs: dict) -> Optional[DocumentQualityRecord]: + """Return the document's quality record, or None if it failed.""" assert _worker_sciencebeam_parser is not None try: - generate_training_data_for_source_filename( + return generate_training_data_for_source_filename( kwargs['source_filename'], output_path=kwargs['output_path'], sciencebeam_parser=_worker_sciencebeam_parser, @@ -1464,10 +1584,66 @@ def _worker_process(kwargs: dict) -> bool: xml_file_list=kwargs['xml_file_list'], enabled_models=kwargs['enabled_models'], ) - return True except Exception: # pylint: disable=broad-except LOGGER.exception('Failed to process %r', kwargs['source_filename']) - return False + return None + + +class _WorkerResult(NamedTuple): + record: Optional[DocumentQualityRecord] + status: str + + @staticmethod + def for_worker_return( + record: Optional[DocumentQualityRecord] + ) -> '_WorkerResult': + """A worker that returned nothing failed; anything else returned its record.""" + return _WorkerResult( + record, + DocumentStatus.OK if record is not None else DocumentStatus.ERROR + ) + + @property + def ok(self) -> bool: + return self.record is not None + + +def _get_worker_result( + async_result, + source_filename: str, + document_timeout: int, +) -> _WorkerResult: + try: + return _WorkerResult.for_worker_return( + async_result.get(timeout=document_timeout or None) + ) + except multiprocessing.TimeoutError: + LOGGER.warning( + 'Document exceeded %ds timeout, skipping: %r', + document_timeout, source_filename, + ) + return _WorkerResult(None, DocumentStatus.TIMEOUT) + except Exception: # pylint: disable=broad-except + LOGGER.exception('Failed to process %r', source_filename) + return _WorkerResult(None, DocumentStatus.ERROR) + + +def _write_quality_record( + quality_writer: Optional[QualityRecordWriter], + source_filename: str, + worker_result: _WorkerResult, +) -> None: + """Record the document, standing in a failed record when the worker returned none.""" + if quality_writer is None: + return + record = worker_result.record + if record is None: + record = get_failed_document_quality_record( + source_filename=source_filename, + document_id=os.path.splitext(os.path.basename(source_filename))[0], + status=worker_result.status, + ) + quality_writer.write(record) def _run_serial( @@ -1477,6 +1653,7 @@ def _run_serial( xml_file_list: Optional[Sequence[str]], progress: '_Progress', document_timeout: int = 0, + quality_writer: Optional[QualityRecordWriter] = None, ) -> None: """Run documents sequentially. @@ -1504,8 +1681,11 @@ def _run_serial( for source_filename in source_file_list: kwargs = {'source_filename': source_filename, **common_kwargs} t0 = time.monotonic() - ok = _worker_process(kwargs) - progress.record(source_filename, ok=ok, elapsed_s=time.monotonic() - t0) + worker_result = _WorkerResult.for_worker_return(_worker_process(kwargs)) + _write_quality_record(quality_writer, source_filename, worker_result) + progress.record( + source_filename, ok=worker_result.ok, elapsed_s=time.monotonic() - t0 + ) return pool = multiprocessing.Pool(1, initializer=_worker_init) # pylint: disable=consider-using-with @@ -1514,22 +1694,19 @@ def _run_serial( kwargs = {'source_filename': source_filename, **common_kwargs} t0 = time.monotonic() async_result = pool.apply_async(_worker_process, (kwargs,)) - try: - ok = async_result.get(timeout=document_timeout) - except multiprocessing.TimeoutError: - LOGGER.warning( - 'Document exceeded %ds timeout, skipping: %r', - document_timeout, source_filename, - ) + worker_result = _get_worker_result( + async_result, source_filename, document_timeout + ) + if worker_result.status == DocumentStatus.TIMEOUT: + # The worker may still be stuck inside a C extension. pool.terminate() pool.join() # pylint: disable-next=consider-using-with pool = multiprocessing.Pool(1, initializer=_worker_init) - ok = False - except Exception: # pylint: disable=broad-except - LOGGER.exception('Failed to process %r', source_filename) - ok = False - progress.record(source_filename, ok=ok, elapsed_s=time.monotonic() - t0) + _write_quality_record(quality_writer, source_filename, worker_result) + progress.record( + source_filename, ok=worker_result.ok, elapsed_s=time.monotonic() - t0 + ) finally: pool.close() pool.join() @@ -1543,6 +1720,7 @@ def _run_parallel_workers( progress: '_Progress', num_workers: int, document_timeout: int = 0, + quality_writer: Optional[QualityRecordWriter] = None, ) -> None: """Process documents in parallel using a multiprocessing.Pool. @@ -1560,7 +1738,6 @@ def _run_parallel_workers( } # pylint: disable-next=consider-using-with pool = multiprocessing.Pool(num_workers, initializer=_worker_init) - timeout_arg = document_timeout if document_timeout > 0 else None work = [ (sf, pool.apply_async(_worker_process, ({'source_filename': sf, **common_kwargs},))) for sf in source_file_list @@ -1568,18 +1745,13 @@ def _run_parallel_workers( try: for source_filename, async_result in work: t0 = time.monotonic() - try: - ok = async_result.get(timeout=timeout_arg) - except multiprocessing.TimeoutError: - LOGGER.warning( - 'Document exceeded %ds timeout, skipping: %r', - document_timeout, source_filename, - ) - ok = False - except Exception: # pylint: disable=broad-except - LOGGER.exception('Failed to process %r', source_filename) - ok = False - progress.record(source_filename, ok=ok, elapsed_s=time.monotonic() - t0) + worker_result = _get_worker_result( + async_result, source_filename, document_timeout + ) + _write_quality_record(quality_writer, source_filename, worker_result) + progress.record( + source_filename, ok=worker_result.ok, elapsed_s=time.monotonic() - t0 + ) finally: pool.terminate() pool.join() @@ -1605,16 +1777,24 @@ def run(args: argparse.Namespace): num_workers = getattr(args, 'num_workers', 1) document_timeout: int = getattr(args, 'document_timeout', 0) - if num_workers > 1: - _run_parallel_workers( - source_file_list, output_path, args, xml_file_list, progress, num_workers, - document_timeout=document_timeout, - ) - else: - _run_serial( - source_file_list, output_path, args, xml_file_list, progress, - document_timeout=document_timeout, - ) + with QualityRecordWriter( + output_path, + model_names=get_enabled_model_names(args.enabled_models), + use_directory_structure=args.use_directory_structure, + ) as quality_writer: + if num_workers > 1: + _run_parallel_workers( + source_file_list, output_path, args, xml_file_list, progress, num_workers, + document_timeout=document_timeout, + quality_writer=quality_writer, + ) + else: + _run_serial( + source_file_list, output_path, args, xml_file_list, progress, + document_timeout=document_timeout, + quality_writer=quality_writer, + ) + LOGGER.info('quality records written: %d', quality_writer.written_count) if progress.n_err: LOGGER.warning('%d/%d documents failed', progress.n_err, total) diff --git a/sciencebeam_parser/training/jats/annotated_document.py b/sciencebeam_parser/training/jats/annotated_document.py index e6680cbf..ae66968b 100644 --- a/sciencebeam_parser/training/jats/annotated_document.py +++ b/sciencebeam_parser/training/jats/annotated_document.py @@ -34,6 +34,19 @@ def set_token_label( ) -> None: self.token_label_by_id[id(token)] = (field_name, sub_field_name, instance_id) + def get_aligned_instance_count(self, field_name: str) -> int: + """Number of instances of a repeated field that got at least one token. + + Instance ids are assigned per parent match and start at 1, so this is the + count of values the aligner placed -- fewer than the JATS holds means + alignment lost them, rather than a later stage. + """ + return len({ + entry[2] + for entry in self.token_label_by_id.values() + if entry[0] == field_name and entry[2] + }) + def coverage_ratio(self) -> float: total = sum(1 for _ in self.layout_document.iter_all_tokens()) if total == 0: diff --git a/sciencebeam_parser/training/jats/field_extractor.py b/sciencebeam_parser/training/jats/field_extractor.py index 5f4430b6..30e84175 100644 --- a/sciencebeam_parser/training/jats/field_extractor.py +++ b/sciencebeam_parser/training/jats/field_extractor.py @@ -1,4 +1,5 @@ -from typing import Dict, Iterator, List, Optional, Sequence, Tuple +from itertools import chain +from typing import Dict, FrozenSet, Iterator, List, Optional, Sequence, Tuple from dataclasses import dataclass from lxml import etree @@ -132,6 +133,30 @@ def _iter_reference_author_values( ) +def iter_reference_sub_field_names( + root: etree._Element, +) -> Iterator[FrozenSet[str]]: + """Yield the sub-field names the JATS carries for each reference, in JATS order. + + The reference set is the one iter_field_values emits a REFERENCE value for, + so a `<ref-list>` inside a `<sub-article>` is not in it: those entries are + never offered to the aligner and cannot be a shortfall. + """ + for ref_el in root.xpath('back/ref-list/ref'): + if not _reference_parent_text(ref_el): + continue + yield frozenset( + field_value.sub_field_name + for field_value in chain( + _iter_reference_author_values(ref_el, JatsFieldNames.REFERENCE), + _iter_sub_field_values( + ref_el, JatsFieldNames.REFERENCE, _REFERENCE_SUB_FIELDS + ), + ) + if field_value.sub_field_name + ) + + def _local_tag(el: etree._Element) -> str: tag = el.tag if isinstance(tag, str) and tag.startswith('{'): diff --git a/sciencebeam_parser/training/quality/__init__.py b/sciencebeam_parser/training/quality/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/sciencebeam_parser/training/quality/counting.py b/sciencebeam_parser/training/quality/counting.py new file mode 100644 index 00000000..f74dbe15 --- /dev/null +++ b/sciencebeam_parser/training/quality/counting.py @@ -0,0 +1,110 @@ +"""Counting the cardinality of generated training data at each stage. + +The reference of record is what `JatsFieldExtractor` emits, and every count is +taken per document so that a corpus total cannot hide a document that lost its +references. +""" +import logging +from typing import Dict, FrozenSet, Iterable, Mapping, Optional, Sequence + +from lxml import etree + +from sciencebeam_parser.models.data import LayoutModelData +from sciencebeam_parser.training.jats.field_vocab import CITATION_LABEL_BY_SUB_FIELD + + +LOGGER = logging.getLogger(__name__) + + +# The element a model writes once per entity, by model name. Counting elements +# only means something for a model whose labels mark repeated entities. +ENTITY_ELEMENT_NAME_BY_MODEL: Mapping[str, str] = { + 'reference-segmenter': 'bibl', + 'citation': 'bibl', +} + +# Models whose labels mark regions rather than repeated entities: an element +# count says nothing about them, and they are listed rather than left out so +# that a model with no entry is a missing decision rather than a silent pass. +MODELS_WITHOUT_ENTITY_COUNT: FrozenSet[str] = frozenset({ + 'segmentation', + 'header', + 'affiliation-address', + 'name-header', + 'name-citation', + 'fulltext', + 'figure', + 'table', +}) + + +def get_local_name(element: etree._Element) -> str: + tag = element.tag + if not isinstance(tag, str): + return '' + return tag.split('}', 1)[1] if tag.startswith('{') else tag + + +def count_entity_elements(model_name: str, tei_root: etree._Element) -> Optional[int]: + """Return the number of entity elements the TEI holds, or None if the model has no count.""" + element_name = ENTITY_ELEMENT_NAME_BY_MODEL.get(model_name) + if element_name is None: + if model_name not in MODELS_WITHOUT_ENTITY_COUNT: + LOGGER.warning( + 'no cardinality check defined for model %r, and it is not declared as having none', + model_name + ) + return None + return sum( + 1 for element in tei_root.iter() if get_local_name(element) == element_name + ) + + +def get_labels_for_model_data_list( + model_data_list: Sequence[LayoutModelData] +) -> FrozenSet[str]: + """Return the labels marked in one entity's model data, without the B- prefix. + + Unlabeled model data carries no label attribute at all, which reads here as + an entity that marks nothing. + """ + labels = ( + getattr(model_data, 'label', None) + for model_data in model_data_list + ) + return frozenset( + label[2:] if label.startswith('B-') else label + for label in labels + if label + ) + + +def count_citation_labels( + reference_sub_field_names: Iterable[FrozenSet[str]], + model_data_list_list: Sequence[Sequence[LayoutModelData]], +) -> Dict[str, Dict[str, int]]: + """Count, per citation label, references whose JATS carries it and references marking it. + + Presence per reference rather than occurrences, since the label conventions + differ: `<author>` covers a whole author list where `<pages>` is written once + per page number. The two sides are counted independently -- a reference + legitimately does not print everything its JATS carries, so what the record + holds is a rate to compare across regenerations, not a per-reference + requirement. + """ + counts: Dict[str, Dict[str, int]] = {} + + def _entry(label: str) -> Dict[str, int]: + return counts.setdefault(label, {'jats': 0, 'marked': 0}) + + for sub_field_names in reference_sub_field_names: + for label in { + CITATION_LABEL_BY_SUB_FIELD[sub_field_name] + for sub_field_name in sub_field_names + if sub_field_name in CITATION_LABEL_BY_SUB_FIELD + }: + _entry(label)['jats'] += 1 + for model_data_list in model_data_list_list: + for label in get_labels_for_model_data_list(model_data_list): + _entry(label)['marked'] += 1 + return counts diff --git a/sciencebeam_parser/training/quality/record.py b/sciencebeam_parser/training/quality/record.py new file mode 100644 index 00000000..4cdf8c55 --- /dev/null +++ b/sciencebeam_parser/training/quality/record.py @@ -0,0 +1,187 @@ +"""The per-document quality record a generated corpus carries. + +One file per model, one row per source document, written as the run proceeds and +whether the document succeeded or not: a document that produced no output at all +is invisible to anything that iterates generated files, and is the case this +record exists to make visible. + +The record is per model because generation is run per model -- a corpus holds one +model's data at one document set and another model's at a different one, so a +record covering a whole corpus would describe the last run rather than the data +beside it. +""" +import json +import logging +import os +from dataclasses import dataclass, field +from typing import IO, Any, Dict, Optional, Sequence + +from sciencebeam_parser.utils.io import auto_uploading_output_file + + +LOGGER = logging.getLogger(__name__) + + +QUALITY_RECORD_FILENAME = 'quality.jsonl' + + +class DocumentStatus: + OK = 'ok' + ERROR = 'error' + TIMEOUT = 'timeout' + + +class JatsStatus: + OK = 'ok' + MISSING = 'missing' + UNPARSABLE = 'unparsable' + UNREADABLE = 'unreadable' + + +def get_quality_record_file_path( + output_path: str, + model_name: str, + use_directory_structure: bool +) -> str: + """Where a model's record goes, following the layout of its training data.""" + if use_directory_structure: + return os.path.join(output_path, model_name, QUALITY_RECORD_FILENAME) + return os.path.join(output_path, model_name + '.' + QUALITY_RECORD_FILENAME) + + +@dataclass +class JatsQualityRecord: + status: str + reference_count: Optional[int] = None + aligned_reference_count: Optional[int] = None + + def to_json_dict(self) -> Dict[str, Any]: + json_dict: Dict[str, Any] = {'status': self.status} + if self.reference_count is not None: + json_dict['reference_count'] = self.reference_count + if self.aligned_reference_count is not None: + json_dict['aligned_reference_count'] = self.aligned_reference_count + return json_dict + + +@dataclass +class ModelQualityRecord: + """What one model wrote for one document. + + `entity_element_count` is None for a model whose labels mark regions rather + than repeated entities; `written` is False when the generator produced no + entities and so wrote no file. + """ + model_name: str + written: bool + entity_element_count: Optional[int] = None + label_counts: Optional[Dict[str, Dict[str, int]]] = None + + def to_json_dict(self) -> Dict[str, Any]: + json_dict: Dict[str, Any] = {'written': self.written} + if self.entity_element_count is not None: + json_dict['entity_element_count'] = self.entity_element_count + if self.label_counts: + json_dict['label_counts'] = self.label_counts + return json_dict + + +@dataclass +class DocumentQualityRecord: + """Everything measured for one document, across the models a run generated for.""" + document_id: str + source_filename: str + status: str = DocumentStatus.OK + jats: JatsQualityRecord = field( + default_factory=lambda: JatsQualityRecord(status=JatsStatus.MISSING) + ) + models: Sequence[ModelQualityRecord] = () + + def to_json_dict_by_model( + self, + model_names: Sequence[str] + ) -> Dict[str, Dict[str, Any]]: + """One row per model, so that each model's record covers the whole document set. + + A document that failed carries no model record, and still gets a row for + every model the run was asked for: it is absent from all of their data. + """ + model_record_by_name = { + model_record.model_name: model_record + for model_record in self.models + } + return { + model_name: { + 'document_id': self.document_id, + 'source_filename': self.source_filename, + 'status': self.status, + 'model': model_name, + 'jats': self.jats.to_json_dict(), + **( + model_record_by_name[model_name].to_json_dict() + if model_name in model_record_by_name + else {} + ), + } + for model_name in model_names + } + + +def get_failed_document_quality_record( + source_filename: str, + document_id: str, + status: str, +) -> DocumentQualityRecord: + return DocumentQualityRecord( + document_id=document_id, + source_filename=source_filename, + status=status, + ) + + +class QualityRecordWriter: + """Append records as JSON lines, one file per model, flushing each line. + + The writer runs in the parent process, which knows the document set the run + was asked for -- a worker that timed out or died cannot report itself. + """ + def __init__( + self, + output_path: str, + model_names: Sequence[str], + use_directory_structure: bool = True + ): + self.output_path = output_path + self.model_names = list(model_names) + self.use_directory_structure = use_directory_structure + self._file_context_by_model: Dict[str, Any] = {} + self._file_by_model: Dict[str, IO] = {} + self.written_count = 0 + + def __enter__(self) -> 'QualityRecordWriter': + for model_name in self.model_names: + file_path = get_quality_record_file_path( + self.output_path, model_name, self.use_directory_structure + ) + LOGGER.info('writing quality record to: %r', file_path) + file_context = auto_uploading_output_file( + file_path, mode='w', encoding='utf-8' + ) + self._file_context_by_model[model_name] = file_context + self._file_by_model[model_name] = file_context.__enter__() + return self + + def __exit__(self, *args) -> None: + for file_context in self._file_context_by_model.values(): + file_context.__exit__(*args) + self._file_context_by_model.clear() + self._file_by_model.clear() + + def write(self, record: DocumentQualityRecord) -> None: + for model_name, json_dict in record.to_json_dict_by_model( + self.model_names + ).items(): + output_file = self._file_by_model[model_name] + output_file.write(json.dumps(json_dict, sort_keys=True) + '\n') + output_file.flush() + self.written_count += 1 diff --git a/tests/training/cli/generate_data_test.py b/tests/training/cli/generate_data_test.py index 34343806..617eb2e9 100644 --- a/tests/training/cli/generate_data_test.py +++ b/tests/training/cli/generate_data_test.py @@ -61,6 +61,7 @@ generate_training_data_for_layout_document, main, ) +from sciencebeam_parser.training.quality.record import DocumentStatus, JatsStatus from tests.processors.fulltext.model_mocks import MockFullTextModels from tests.test_utils import log_on_exception @@ -426,6 +427,179 @@ def test_should_generate_data_using_mock_models( # noqa pylint: disable=too-man tei_expected_values=[sample_layout_document.ref_title_block.text] ) + def test_should_return_a_quality_record_with_the_count_per_model( + self, + tmp_path: Path, + sample_layout_document: SampleLayoutDocument, + fulltext_models_mock: MockFullTextModels + ): + configure_fulltext_models_mock_with_sample_document( + fulltext_models_mock, + sample_layout_document + ) + + output_path = tmp_path / 'output' + output_path.mkdir() + record = generate_training_data_for_layout_document( + layout_document=sample_layout_document.layout_document, + output_path=str(output_path), + source_filename=SOURCE_FILENAME_1, + document_features_context=DEFAULT_DOCUMENT_FEATURES_CONTEXT, + fulltext_models=fulltext_models_mock, + use_model=True, + use_directory_structure=False + ) + + json_dict_by_model = record.to_json_dict_by_model( + ['reference-segmenter', 'citation', 'segmentation'] + ) + reference_segmenter_json_dict = json_dict_by_model['reference-segmenter'] + assert reference_segmenter_json_dict['document_id'] == 'test1' + assert reference_segmenter_json_dict['status'] == DocumentStatus.OK + assert reference_segmenter_json_dict['written'] is True + assert reference_segmenter_json_dict['entity_element_count'] == 1 + assert json_dict_by_model['citation']['entity_element_count'] == 1 + assert 'entity_element_count' not in json_dict_by_model['segmentation'] + + def test_should_report_a_missing_jats_without_a_reference_count( + self, + tmp_path: Path, + sample_layout_document: SampleLayoutDocument, + fulltext_models_mock: MockFullTextModels + ): + configure_fulltext_models_mock_with_sample_document( + fulltext_models_mock, + sample_layout_document + ) + + output_path = tmp_path / 'output' + output_path.mkdir() + record = generate_training_data_for_layout_document( + layout_document=sample_layout_document.layout_document, + output_path=str(output_path), + source_filename=SOURCE_FILENAME_1, + document_features_context=DEFAULT_DOCUMENT_FEATURES_CONTEXT, + fulltext_models=fulltext_models_mock, + use_model=True, + use_directory_structure=False + ) + + assert record.to_json_dict_by_model( + ['reference-segmenter'] + )['reference-segmenter']['jats'] == {'status': JatsStatus.MISSING} + + def test_should_report_a_jats_that_could_not_be_parsed( + self, + tmp_path: Path, + sample_layout_document: SampleLayoutDocument, + fulltext_models_mock: MockFullTextModels + ): + configure_fulltext_models_mock_with_sample_document( + fulltext_models_mock, + sample_layout_document + ) + empty_jats_path = tmp_path / 'empty.jats.xml' + empty_jats_path.write_bytes(b'') + + output_path = tmp_path / 'output' + output_path.mkdir() + record = generate_training_data_for_layout_document( + layout_document=sample_layout_document.layout_document, + output_path=str(output_path), + source_filename=SOURCE_FILENAME_1, + document_features_context=DEFAULT_DOCUMENT_FEATURES_CONTEXT, + fulltext_models=fulltext_models_mock, + use_model=True, + use_directory_structure=False, + jats_xml_filename=str(empty_jats_path) + ) + + assert record.to_json_dict_by_model( + ['reference-segmenter'] + )['reference-segmenter']['jats'] == {'status': JatsStatus.UNPARSABLE} + + def test_should_report_a_jats_declaring_no_references_as_a_zero_count( + self, + tmp_path: Path, + sample_layout_document: SampleLayoutDocument, + fulltext_models_mock: MockFullTextModels + ): + configure_fulltext_models_mock_with_sample_document( + fulltext_models_mock, + sample_layout_document + ) + jats_path = tmp_path / 'no-references.jats.xml' + jats_path.write_text( + '<article><back><sec><p>Appendix text.</p></sec></back></article>', + encoding='utf-8' + ) + + output_path = tmp_path / 'output' + output_path.mkdir() + record = generate_training_data_for_layout_document( + layout_document=sample_layout_document.layout_document, + output_path=str(output_path), + source_filename=SOURCE_FILENAME_1, + document_features_context=DEFAULT_DOCUMENT_FEATURES_CONTEXT, + fulltext_models=fulltext_models_mock, + use_model=True, + use_directory_structure=False, + jats_xml_filename=str(jats_path) + ) + + assert record.to_json_dict_by_model( + ['reference-segmenter'] + )['reference-segmenter']['jats'] == { + 'status': JatsStatus.OK, + 'reference_count': 0, + 'aligned_reference_count': 0 + } + + def test_should_count_the_jats_references_and_the_citation_labels( + self, + tmp_path: Path, + sample_layout_document: SampleLayoutDocument, + fulltext_models_mock: MockFullTextModels + ): + configure_fulltext_models_mock_with_sample_document( + fulltext_models_mock, + sample_layout_document + ) + jats_path = tmp_path / 'references.jats.xml' + jats_path.write_text( + '<article><back><ref-list>' + '<ref><label>1</label><element-citation>' + '<article-title>Reference 1</article-title>' + '<person-group person-group-type="author">' + '<name><surname>Ref Author Surname 1</surname></name>' + '</person-group>' + '<pub-id pub-id-type="doi">10.1234/not-printed</pub-id>' + '</element-citation></ref>' + '</ref-list></back></article>', + encoding='utf-8' + ) + + output_path = tmp_path / 'output' + output_path.mkdir() + record = generate_training_data_for_layout_document( + layout_document=sample_layout_document.layout_document, + output_path=str(output_path), + source_filename=SOURCE_FILENAME_1, + document_features_context=DEFAULT_DOCUMENT_FEATURES_CONTEXT, + fulltext_models=fulltext_models_mock, + use_model=True, + use_directory_structure=False, + jats_xml_filename=str(jats_path) + ) + + json_dict = record.to_json_dict_by_model(['citation'])['citation'] + assert json_dict['jats']['reference_count'] == 1 + assert json_dict['jats']['aligned_reference_count'] == 1 + label_counts = json_dict['label_counts'] + assert label_counts['<title>'] == {'jats': 1, 'marked': 1} + # The DOI is in the JATS and not on the page: counted, and not marked. + assert label_counts[IDENTIFIER_LABEL] == {'jats': 1, 'marked': 0} + def test_not_should_generate_figure_data_if_not_present( # noqa pylint: disable=too-many-locals, too-many-statements self, tmp_path: Path, diff --git a/tests/training/jats/test_field_extractor.py b/tests/training/jats/test_field_extractor.py index bd403a32..42568fa0 100644 --- a/tests/training/jats/test_field_extractor.py +++ b/tests/training/jats/test_field_extractor.py @@ -2,7 +2,10 @@ from lxml import etree -from sciencebeam_parser.training.jats.field_extractor import JatsFieldExtractor +from sciencebeam_parser.training.jats.field_extractor import ( + JatsFieldExtractor, + iter_reference_sub_field_names +) from sciencebeam_parser.training.jats.field_vocab import JatsFieldNames, JatsSubFieldNames @@ -469,3 +472,66 @@ def test_main_article_body_not_labeled_as_sub_article(self): assert 'Review text' in sub[0].text assert len(body) == 1 assert 'Main article' in body[0].text + + +class TestIterReferenceSubFieldNames: + def test_yields_the_sub_fields_each_reference_carries(self): + sub_field_names = list(iter_reference_sub_field_names(_parse_jats( + '<article><back><ref-list>' + '<ref><label>1</label><element-citation>' + '<person-group person-group-type="author"><name><surname>Smith</surname></name>' + '</person-group>' + '<article-title>First Title</article-title><source>Journal One</source>' + '<year>2020</year>' + '</element-citation></ref>' + '<ref><element-citation>' + '<article-title>Second Title</article-title>' + '<pub-id pub-id-type="doi">10.1234/abc</pub-id>' + '</element-citation></ref>' + '</ref-list></back></article>' + ))) + assert sub_field_names == [ + frozenset({ + JatsSubFieldNames.REFERENCE_LABEL, + JatsSubFieldNames.REFERENCE_AUTHOR, + JatsSubFieldNames.REFERENCE_ARTICLE_TITLE, + JatsSubFieldNames.REFERENCE_SOURCE, + JatsSubFieldNames.REFERENCE_YEAR, + }), + frozenset({ + JatsSubFieldNames.REFERENCE_ARTICLE_TITLE, + JatsSubFieldNames.REFERENCE_DOI, + }), + ] + + def test_excludes_a_sub_article_reference_list(self): + # Those entries are never offered to the aligner, so counting them would + # report a shortfall the pipeline never had. + sub_field_names = list(iter_reference_sub_field_names(_parse_jats( + '<article>' + '<back><ref-list><ref><element-citation>' + '<article-title>Main Title</article-title>' + '</element-citation></ref></ref-list></back>' + '<sub-article article-type="peer-review">' + '<back><ref-list><ref><element-citation>' + '<article-title>Review Title</article-title>' + '</element-citation></ref></ref-list></back>' + '</sub-article>' + '</article>' + ))) + assert sub_field_names == [frozenset({JatsSubFieldNames.REFERENCE_ARTICLE_TITLE})] + + def test_skips_a_reference_with_no_text(self): + sub_field_names = list(iter_reference_sub_field_names(_parse_jats( + '<article><back><ref-list>' + '<ref><element-citation><article-title>A Title</article-title>' + '</element-citation></ref>' + '<ref><element-citation/></ref>' + '</ref-list></back></article>' + ))) + assert sub_field_names == [frozenset({JatsSubFieldNames.REFERENCE_ARTICLE_TITLE})] + + def test_yields_nothing_for_a_jats_declaring_no_references(self): + assert not list(iter_reference_sub_field_names(_parse_jats( + '<article><back><sec><p>Some appendix text.</p></sec></back></article>' + ))) diff --git a/tests/training/quality/__init__.py b/tests/training/quality/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/training/quality/counting_test.py b/tests/training/quality/counting_test.py new file mode 100644 index 00000000..564045ef --- /dev/null +++ b/tests/training/quality/counting_test.py @@ -0,0 +1,122 @@ +from lxml import etree + +from sciencebeam_parser.models.citation.labels import IDENTIFIER_LABEL +from sciencebeam_parser.models.data import LabeledLayoutModelData, LayoutModelData +from sciencebeam_parser.training.jats.field_vocab import JatsSubFieldNames +from sciencebeam_parser.training.quality.counting import ( + count_citation_labels, + count_entity_elements, + get_labels_for_model_data_list +) + + +TEI_NS = 'http://www.tei-c.org/ns/1.0' + + +def _labeled(label) -> LabeledLayoutModelData: + return LabeledLayoutModelData(data_line='token', label=label) + + +class TestCountEntityElements: + def test_should_count_reference_segmenter_elements(self): + tei_root = etree.fromstring( + '<tei><text><listBibl>' + '<bibl>Reference 1</bibl><bibl>Reference 2</bibl>' + '</listBibl></text></tei>' + ) + assert count_entity_elements('reference-segmenter', tei_root) == 2 + + def test_should_count_namespaced_citation_elements(self): + tei_root = etree.fromstring( + f'<TEI xmlns="{TEI_NS}"><text><back><listBibl>' + '<bibl>Reference 1</bibl><bibl>Reference 2</bibl><bibl>Reference 3</bibl>' + '</listBibl></back></text></TEI>' + ) + assert count_entity_elements('citation', tei_root) == 3 + + def test_should_count_an_element_holding_only_a_label(self): + # An element with no token of its own still exists in the TEI, and is the + # difference between the element count and the entity count after parsing. + tei_root = etree.fromstring( + '<tei><text><listBibl>' + '<bibl><label>1</label>Reference 1</bibl><bibl><label>2</label></bibl>' + '</listBibl></text></tei>' + ) + assert count_entity_elements('reference-segmenter', tei_root) == 2 + + def test_should_return_none_for_a_model_declared_as_having_no_entity_count(self): + tei_root = etree.fromstring('<tei><text>Some text</text></tei>') + assert count_entity_elements('segmentation', tei_root) is None + + def test_should_return_none_for_an_undeclared_model(self): + tei_root = etree.fromstring('<tei><text>Some text</text></tei>') + assert count_entity_elements('not-a-model', tei_root) is None + + +class TestGetLabelsForModelDataList: + def test_should_strip_the_begin_prefix(self): + assert get_labels_for_model_data_list([ + _labeled('B-<title>'), _labeled('<title>') + ]) == {'<title>'} + + def test_should_ignore_unlabeled_model_data(self): + assert get_labels_for_model_data_list([ + LayoutModelData(data_line='token'), _labeled(None), _labeled('<date>') + ]) == {'<date>'} + + +class TestCountCitationLabels: + def test_should_count_a_label_the_jats_carries_and_the_data_marks(self): + counts = count_citation_labels( + [frozenset({JatsSubFieldNames.REFERENCE_ARTICLE_TITLE})], + [[_labeled('<title>')]] + ) + assert counts['<title>'] == {'jats': 1, 'marked': 1} + + def test_should_count_a_sub_field_the_page_does_not_print_as_unmarked(self): + counts = count_citation_labels( + [frozenset({JatsSubFieldNames.REFERENCE_DOI})], + [[_labeled('<title>')]] + ) + assert counts[IDENTIFIER_LABEL] == {'jats': 1, 'marked': 0} + + def test_should_count_identifiers_of_different_kinds_under_one_label(self): + counts = count_citation_labels( + [frozenset({ + JatsSubFieldNames.REFERENCE_DOI, JatsSubFieldNames.REFERENCE_PMID + })], + [[_labeled(IDENTIFIER_LABEL), _labeled('B-' + IDENTIFIER_LABEL)]] + ) + assert counts[IDENTIFIER_LABEL] == {'jats': 1, 'marked': 1} + + def test_should_count_presence_per_reference_rather_than_occurrences(self): + # <author> covers a whole author list, so several author tokens in one + # reference are one reference marking the label. + counts = count_citation_labels( + [frozenset({JatsSubFieldNames.REFERENCE_AUTHOR})], + [[_labeled('B-<author>'), _labeled('<author>'), _labeled('<author>')]] + ) + assert counts['<author>'] == {'jats': 1, 'marked': 1} + + def test_should_count_a_sub_field_appearing_twice_in_one_reference_once(self): + # fpage and lpage both map to <pages>, and <pages> is written once per + # page number, so neither side may count an occurrence. + counts = count_citation_labels( + [frozenset({ + JatsSubFieldNames.REFERENCE_FPAGE, JatsSubFieldNames.REFERENCE_LPAGE + })], + [[_labeled('B-<pages>'), _labeled('B-<pages>')]] + ) + assert counts['<pages>'] == {'jats': 1, 'marked': 1} + + def test_should_count_over_all_references(self): + counts = count_citation_labels( + [ + frozenset({JatsSubFieldNames.REFERENCE_ARTICLE_TITLE}), + frozenset({JatsSubFieldNames.REFERENCE_ARTICLE_TITLE}), + frozenset({JatsSubFieldNames.REFERENCE_YEAR}), + ], + [[_labeled('<title>')], [_labeled('<date>')], [_labeled('<date>')]] + ) + assert counts['<title>'] == {'jats': 2, 'marked': 1} + assert counts['<date>'] == {'jats': 1, 'marked': 2} diff --git a/tests/training/quality/record_test.py b/tests/training/quality/record_test.py new file mode 100644 index 00000000..3fa32db1 --- /dev/null +++ b/tests/training/quality/record_test.py @@ -0,0 +1,184 @@ +import json +from pathlib import Path + +from sciencebeam_parser.training.quality.record import ( + DocumentQualityRecord, + DocumentStatus, + JatsQualityRecord, + JatsStatus, + ModelQualityRecord, + QualityRecordWriter, + get_failed_document_quality_record, + get_quality_record_file_path +) + + +DOCUMENT_ID_1 = 'document1' +SOURCE_FILENAME_1 = '/source/document1.pdf' +REFERENCE_SEGMENTER = 'reference-segmenter' + + +def _record_for_models(*models: ModelQualityRecord) -> DocumentQualityRecord: + return DocumentQualityRecord( + document_id=DOCUMENT_ID_1, + source_filename=SOURCE_FILENAME_1, + jats=JatsQualityRecord( + status=JatsStatus.OK, reference_count=45, aligned_reference_count=2 + ), + models=models, + ) + + +class TestGetQualityRecordFilePath: + def test_should_use_a_directory_per_model(self): + assert get_quality_record_file_path( + '/output', REFERENCE_SEGMENTER, use_directory_structure=True + ) == '/output/reference-segmenter/quality.jsonl' + + def test_should_stay_flat_without_the_directory_structure(self): + assert get_quality_record_file_path( + '/output', REFERENCE_SEGMENTER, use_directory_structure=False + ) == '/output/reference-segmenter.quality.jsonl' + + +class TestDocumentQualityRecord: + def test_should_hold_the_count_at_each_stage(self): + json_dict = _record_for_models( + ModelQualityRecord( + model_name=REFERENCE_SEGMENTER, written=True, entity_element_count=2 + ) + ).to_json_dict_by_model([REFERENCE_SEGMENTER])[REFERENCE_SEGMENTER] + assert json_dict['document_id'] == DOCUMENT_ID_1 + assert json_dict['model'] == REFERENCE_SEGMENTER + assert json_dict['status'] == DocumentStatus.OK + assert json_dict['jats'] == { + 'status': JatsStatus.OK, 'reference_count': 45, 'aligned_reference_count': 2 + } + assert json_dict['written'] is True + assert json_dict['entity_element_count'] == 2 + + def test_should_return_one_row_per_model(self): + json_dict_by_model = _record_for_models( + ModelQualityRecord( + model_name=REFERENCE_SEGMENTER, written=True, entity_element_count=2 + ), + ModelQualityRecord( + model_name='citation', + written=True, + entity_element_count=2, + label_counts={'<title>': {'jats': 44, 'marked': 2}} + ), + ).to_json_dict_by_model([REFERENCE_SEGMENTER, 'citation']) + assert set(json_dict_by_model) == {REFERENCE_SEGMENTER, 'citation'} + assert 'label_counts' not in json_dict_by_model[REFERENCE_SEGMENTER] + assert json_dict_by_model['citation']['label_counts'] == { + '<title>': {'jats': 44, 'marked': 2} + } + + def test_should_record_a_model_that_wrote_no_file(self): + json_dict = _record_for_models( + ModelQualityRecord( + model_name=REFERENCE_SEGMENTER, written=False, entity_element_count=0 + ) + ).to_json_dict_by_model([REFERENCE_SEGMENTER])[REFERENCE_SEGMENTER] + assert json_dict['written'] is False + assert json_dict['entity_element_count'] == 0 + + def test_should_omit_an_entity_count_a_model_does_not_have(self): + json_dict = _record_for_models( + ModelQualityRecord(model_name='segmentation', written=True) + ).to_json_dict_by_model(['segmentation'])['segmentation'] + assert json_dict['written'] is True + assert 'entity_element_count' not in json_dict + + def test_should_record_a_jats_that_could_not_be_parsed_without_counts(self): + json_dict = DocumentQualityRecord( + document_id=DOCUMENT_ID_1, + source_filename=SOURCE_FILENAME_1, + jats=JatsQualityRecord(status=JatsStatus.UNPARSABLE), + models=[ModelQualityRecord(model_name=REFERENCE_SEGMENTER, written=False)], + ).to_json_dict_by_model([REFERENCE_SEGMENTER])[REFERENCE_SEGMENTER] + assert json_dict['jats'] == {'status': JatsStatus.UNPARSABLE} + + def test_should_record_a_jats_declaring_no_references_as_a_zero_count(self): + json_dict = DocumentQualityRecord( + document_id=DOCUMENT_ID_1, + source_filename=SOURCE_FILENAME_1, + jats=JatsQualityRecord( + status=JatsStatus.OK, reference_count=0, aligned_reference_count=0 + ), + models=[ModelQualityRecord(model_name=REFERENCE_SEGMENTER, written=False)], + ).to_json_dict_by_model([REFERENCE_SEGMENTER])[REFERENCE_SEGMENTER] + assert json_dict['jats']['reference_count'] == 0 + + +class TestGetFailedDocumentQualityRecord: + def test_should_give_a_timed_out_document_a_row_for_every_model(self): + json_dict_by_model = get_failed_document_quality_record( + source_filename=SOURCE_FILENAME_1, + document_id=DOCUMENT_ID_1, + status=DocumentStatus.TIMEOUT, + ).to_json_dict_by_model([REFERENCE_SEGMENTER, 'citation']) + assert set(json_dict_by_model) == {REFERENCE_SEGMENTER, 'citation'} + for json_dict in json_dict_by_model.values(): + assert json_dict['document_id'] == DOCUMENT_ID_1 + assert json_dict['status'] == DocumentStatus.TIMEOUT + assert 'written' not in json_dict + + +class TestQualityRecordWriter: + def test_should_write_one_file_per_model(self, tmp_path: Path): + with QualityRecordWriter( + str(tmp_path), model_names=[REFERENCE_SEGMENTER, 'citation'] + ) as writer: + writer.write(_record_for_models( + ModelQualityRecord( + model_name=REFERENCE_SEGMENTER, written=True, entity_element_count=2 + ), + ModelQualityRecord( + model_name='citation', written=True, entity_element_count=3 + ), + )) + assert writer.written_count == 1 + for model_name, expected_count in [(REFERENCE_SEGMENTER, 2), ('citation', 3)]: + lines = ( + tmp_path / model_name / 'quality.jsonl' + ).read_text(encoding='utf-8').splitlines() + assert [json.loads(line)['entity_element_count'] for line in lines] == [ + expected_count + ] + + def test_should_write_one_line_per_document(self, tmp_path: Path): + with QualityRecordWriter( + str(tmp_path), model_names=[REFERENCE_SEGMENTER] + ) as writer: + for document_id in ['document1', 'document2']: + writer.write(DocumentQualityRecord( + document_id=document_id, + source_filename=f'/source/{document_id}.pdf', + models=[ + ModelQualityRecord(model_name=REFERENCE_SEGMENTER, written=True) + ], + )) + lines = ( + tmp_path / REFERENCE_SEGMENTER / 'quality.jsonl' + ).read_text(encoding='utf-8').splitlines() + assert [json.loads(line)['document_id'] for line in lines] == [ + 'document1', 'document2' + ] + + def test_should_flush_each_record_so_an_interrupted_run_keeps_what_it_had( + self, tmp_path: Path + ): + with QualityRecordWriter( + str(tmp_path), model_names=[REFERENCE_SEGMENTER] + ) as writer: + writer.write(DocumentQualityRecord( + document_id=DOCUMENT_ID_1, + source_filename=SOURCE_FILENAME_1, + models=[ + ModelQualityRecord(model_name=REFERENCE_SEGMENTER, written=True) + ], + )) + record_file_path = tmp_path / REFERENCE_SEGMENTER / 'quality.jsonl' + assert len(record_file_path.read_text(encoding='utf-8').splitlines()) == 1 From 3675b351deb1b469c1ca311ca48564fc8037ed72 Mon Sep 17 00:00:00 2001 From: Daniel Ecer <d.ecer@elifesciences.org> Date: Thu, 20 Aug 2026 13:45:09 +0100 Subject: [PATCH 2/4] Count the entities training data ends up with, at assembly Generation records the references the JATS declares and the elements it wrote, and stops there: the count that decides what a model trains on is the entities those elements parse back to, and only the delft conversion sees them. Without it a loss is visible but not attributable - scielo_preprints-jats looks 26 elements short of its JATS where alignment loses 47, element writing gains 21 to split references, and the parse loses 5. generate_delft_data now takes that count, joins the record generation wrote via --quality-record-path, writes one row per document to <delft-output-path>.quality.jsonl, and logs a summary per corpus naming the documents that lost entities. It reads the generation record rather than extending it, so generated output stays reproducible from generation alone and can be assembled more than once. The citation model has no entity count to take, since its parser's root element path is bibl and every element is already its own training sequence. It records label starts per sequence instead, which is comparable with what generation recorded as marked, and is where spec 010's defect reached that model. Over both corpora the two sides agree exactly, label by label. Three things this could not have been trusted to do without running it: The two CLIs disagree on model names - generation writes reference-segmenter where this one takes --model-name=reference_segmenter - so every lookup keyed on one spelling returned no counts at all while appearing to work. Names are canonicalised, and the record always carries the hyphenated form so the two sides join. Documents that produced no training data were inferred from the files the glob matched, which reported every document outside a narrowed run: 35 of 40 on a three-file run. It reads the record's own `written` instead. Nothing was printed. The import chain installs a root log handler and raises the root level to ERROR, so basicConfig is a no-op and the summary was visible only under --debug; this CLI now sets its own loggers. Counting only: no threshold, no verdict, and nothing is filtered. --- doc/training.md | 39 +++ .../training/cli/generate_delft_data.py | 173 +++++++++++-- .../training/quality/assembly.py | 199 +++++++++++++++ .../training/quality/counting.py | 82 +++++- .../training/cli/generate_delft_data_test.py | 121 +++++++++ tests/training/quality/assembly_test.py | 237 ++++++++++++++++++ tests/training/quality/counting_test.py | 88 ++++++- 7 files changed, 919 insertions(+), 20 deletions(-) create mode 100644 sciencebeam_parser/training/quality/assembly.py create mode 100644 tests/training/quality/assembly_test.py diff --git a/doc/training.md b/doc/training.md index 68510590..24e9e7b3 100644 --- a/doc/training.md +++ b/doc/training.md @@ -146,6 +146,45 @@ which `delft` models read as a text feature and `wapiti` templates do not reference. Training data generated with the flag cannot be mixed with GROBID's `segmentation` corpus, since it is one column wider. +#### The assembly quality record + +This step is the only place the last count exists: the TEI holds elements, and +what the training data ends up with is the entities those elements parse back to. +It writes `<delft-output-path>.quality.jsonl` (or `--quality-output-path`), one row +per document, and logs a summary per corpus. + +Pass `--quality-record-path` to join what generation recorded, so that a loss can be +attributed to a stage rather than only observed: + +```bash +python -m sciencebeam_parser.training.cli.generate_delft_data \ + --model-name="reference_segmenter" \ + --tei-source-path="data/generated-training-data/train/*/reference-segmenter/corpus/tei/*.tei.xml" \ + --quality-record-path="data/generated-training-data/train/*/reference-segmenter/quality.jsonl" \ + --delft-output-path="./data/generated-training-data/delft/reference-segmenter/corpus/reference-segmenter.data" +``` + +```json +{ + "document_id": "PPR459453", "model": "reference-segmenter", "corpus": "scielo_preprints-jats", + "sequence_count": 1, "entity_start_count": 2, + "generated": {"jats": {"reference_count": 45}, "entity_element_count": 2} +} +``` + +- `entity_start_count` against the generated `entity_element_count` is the parse: fewer + entities than elements is a boundary lost between siblings, and the summary names + the documents it happened to. +- `sequence_count` is the training sequences the document contributes. For `citation` + every element is its own sequence, so it has no entity count and carries + `label_start_counts` instead — per label, the sequences marking it, comparable with + what generation recorded as marked. +- A document generation recorded that produced no training data at all is reported + by id: nothing that iterates the TEI can see it. + +Without `--quality-record-path` the counts are still recorded, with nothing to +compare them against — which is enough to re-check a committed corpus offline. + #### Example command for `segmentation` model ```bash diff --git a/sciencebeam_parser/training/cli/generate_delft_data.py b/sciencebeam_parser/training/cli/generate_delft_data.py index e4dc102e..84defe58 100644 --- a/sciencebeam_parser/training/cli/generate_delft_data.py +++ b/sciencebeam_parser/training/cli/generate_delft_data.py @@ -1,7 +1,7 @@ import argparse import logging import os -from typing import Iterable, List, Optional, Sequence, Tuple +from typing import Iterable, List, Mapping, NamedTuple, Optional, Sequence, Tuple from lxml import etree @@ -39,6 +39,20 @@ select_feature_columns ) +from sciencebeam_parser.training.quality.assembly import ( + AssembledDocumentRecord, + GeneratedDocumentRecord, + get_assembly_summary_by_corpus, + get_document_ids_without_generated_output, + read_generated_document_records, + write_assembly_records +) +from sciencebeam_parser.training.quality.counting import ( + count_entity_starts, + count_label_starts_per_sequence, + get_canonical_model_name, + is_model_counted_by_label +) from sciencebeam_parser.resources.default_config import DEFAULT_CONFIG_FILE from sciencebeam_parser.config.config import AppConfig from sciencebeam_parser.app.parser import ScienceBeamParser @@ -71,6 +85,26 @@ def parse_args(argv: Optional[List[str]] = None) -> argparse.Namespace: type=str, required=True ) + parser.add_argument( + '--quality-record-path', + type=str, + required=False, + help=( + 'File pattern of the quality.jsonl written at generation, e.g.' + ' "<data>/train/*/reference-segmenter/quality.jsonl". Its counts are joined' + ' with the entity count only this step can take. Without it the entity count' + ' is still recorded, with nothing to compare it against.' + ) + ) + parser.add_argument( + '--quality-output-path', + type=str, + required=False, + help=( + 'Where to write the assembly quality record' + ' (default: the delft output path with ".quality.jsonl" appended).' + ) + ) parser.add_argument( '--include-extra-columns', action='store_true', @@ -204,14 +238,19 @@ def get_data_generator_for_model_name( ) -def iter_generate_delft_training_data_lines_for_document( # pylint: disable=too-many-locals +class DelftDocumentResult(NamedTuple): + data_lines: Sequence[str] + labeled_layout_tokens_list: Sequence[Sequence[LabeledLayoutToken]] + + +def get_delft_training_data_for_document( # pylint: disable=too-many-locals tei_file: str, raw_file: Optional[str], training_tei_parser: TrainingTeiParser, data_generator: ModelDataGenerator, column_layout: GrobidColumnLayout, include_extra_columns: bool = False -) -> Iterable[str]: +) -> DelftDocumentResult: with auto_download_input_file( tei_file, auto_decompress=True @@ -258,7 +297,7 @@ def iter_generate_delft_training_data_lines_for_document( # pylint: disable=too _texts, features = load_data_crf_lines(data_line_iterable) LOGGER.debug('features: %r', features) if not len(features): # pylint: disable=len-as-condition - return + return DelftDocumentResult([], labeled_layout_tokens_list) feature_indices = get_validated_training_data_feature_indices( column_layout, feature_column_count=len(features[0][0]), @@ -266,12 +305,88 @@ def iter_generate_delft_training_data_lines_for_document( # pylint: disable=too data_generator_column_names=data_generator.feature_names, include_extra_columns=include_extra_columns ) - yield from iter_format_tag_result( - tag_result=translated_tag_result, - output_format=TagOutputFormats.DATA, - texts=None, - features=select_feature_columns(features, feature_indices) + return DelftDocumentResult( + list(iter_format_tag_result( + tag_result=translated_tag_result, + output_format=TagOutputFormats.DATA, + texts=None, + features=select_feature_columns(features, feature_indices) + )), + labeled_layout_tokens_list + ) + + +def get_document_id_for_tei_file( + tei_file: str, + tei_filename_suffix: Optional[str] +) -> str: + """The document id generation recorded, which is the source name. + + A model with no declared suffix, or a file that does not carry it, falls back + to everything before the first dot. + """ + basename = os.path.basename(tei_file) + if basename.endswith('.gz'): + basename = basename[:-len('.gz')] + if tei_filename_suffix and basename.endswith(tei_filename_suffix): + return basename[:-len(tei_filename_suffix)] + return basename.split('.', maxsplit=1)[0] + + +def get_tei_filename_suffix_for_model_name( + model_name: str, + sciencebeam_parser: ScienceBeamParser +) -> Optional[str]: + model = sciencebeam_parser.fulltext_models.get_sequence_model_by_name(model_name) + return model.get_tei_training_data_generator().get_default_tei_filename_suffix() + + +def get_assembled_document_record( + document_id: str, + model_name: str, + result: DelftDocumentResult, + generated_record_by_document_id: Mapping[str, GeneratedDocumentRecord], +) -> AssembledDocumentRecord: + generated = generated_record_by_document_id.get(document_id) + return AssembledDocumentRecord( + document_id=document_id, + model_name=get_canonical_model_name(model_name), + corpus=generated.corpus if generated else None, + sequence_count=len(result.labeled_layout_tokens_list), + entity_start_count=count_entity_starts( + model_name, result.labeled_layout_tokens_list + ), + label_start_counts=( + count_label_starts_per_sequence(result.labeled_layout_tokens_list) + if is_model_counted_by_label(model_name) + else None + ), + generated=generated, + ) + + +def log_assembly_summary( + model_name: str, + assembled_records: Sequence[AssembledDocumentRecord], + generated_record_by_document_id: Mapping[str, GeneratedDocumentRecord], +) -> None: + canonical_model_name = get_canonical_model_name(model_name) + for corpus, summary in sorted( + get_assembly_summary_by_corpus(assembled_records).items(), + key=lambda item: item[0] or '' + ): + LOGGER.info( + '%s / %s: %s', corpus or 'corpus not known', canonical_model_name, summary + ) + document_ids_without_output = get_document_ids_without_generated_output( + generated_record_by_document_id ) + if document_ids_without_output: + LOGGER.warning( + '%d documents generation wrote no %s file for: %r', + len(document_ids_without_output), canonical_model_name, + document_ids_without_output + ) def generate_delft_training_data( # pylint: disable=too-many-locals @@ -280,7 +395,9 @@ def generate_delft_training_data( # pylint: disable=too-many-locals raw_source_path: str, delft_output_path: str, sciencebeam_parser: ScienceBeamParser, - include_extra_columns: bool = False + include_extra_columns: bool = False, + quality_record_path: Optional[str] = None, + quality_output_path: Optional[str] = None ): training_tei_parser = get_training_tei_parser_for_model_name( model_name, @@ -309,6 +426,15 @@ def generate_delft_training_data( # pylint: disable=too-many-locals else: raw_file_list = [None] * len(tei_file_list) LOGGER.info('raw_file_list: %r', raw_file_list) + generated_record_by_document_id = ( + read_generated_document_records(quality_record_path) + if quality_record_path + else {} + ) + tei_filename_suffix = get_tei_filename_suffix_for_model_name( + model_name, sciencebeam_parser=sciencebeam_parser + ) + assembled_records: List[AssembledDocumentRecord] = [] LOGGER.info('writing to : %r', delft_output_path) with auto_uploading_output_file( delft_output_path, @@ -318,14 +444,28 @@ def generate_delft_training_data( # pylint: disable=too-many-locals for document_index, (tei_file, raw_file) in enumerate(zip(tei_file_list, raw_file_list)): if document_index > 0: data_fp.write('\n\n') - data_fp.writelines(iter_generate_delft_training_data_lines_for_document( + result = get_delft_training_data_for_document( tei_file=tei_file, raw_file=raw_file, training_tei_parser=training_tei_parser, data_generator=data_generator, column_layout=column_layout, include_extra_columns=include_extra_columns + ) + data_fp.writelines(result.data_lines) + assembled_records.append(get_assembled_document_record( + document_id=get_document_id_for_tei_file(tei_file, tei_filename_suffix), + model_name=model_name, + result=result, + generated_record_by_document_id=generated_record_by_document_id, )) + write_assembly_records( + quality_output_path or delft_output_path + '.quality.jsonl', + assembled_records + ) + log_assembly_summary( + model_name, assembled_records, generated_record_by_document_id + ) def run(args: argparse.Namespace): @@ -340,16 +480,21 @@ def run(args: argparse.Namespace): raw_source_path=args.raw_source_path, delft_output_path=args.delft_output_path, sciencebeam_parser=sciencebeam_parser, - include_extra_columns=args.include_extra_columns + include_extra_columns=args.include_extra_columns, + quality_record_path=args.quality_record_path, + quality_output_path=args.quality_output_path ) def main(argv: Optional[List[str]] = None): LOGGER.debug('argv: %r', argv) args = parse_args(argv) + # The import chain installs a root handler and raises the root level, so this + # CLI's own output -- the quality summary included -- is otherwise dropped. + for name in [__name__, 'sciencebeam_parser']: + logging.getLogger(name).setLevel('DEBUG' if args.debug else 'INFO') if args.debug: - for name in [__name__, 'sciencebeam_parser', 'sciencebeam_trainer_delft']: - logging.getLogger(name).setLevel('DEBUG') + logging.getLogger('sciencebeam_trainer_delft').setLevel('DEBUG') run(args) diff --git a/sciencebeam_parser/training/quality/assembly.py b/sciencebeam_parser/training/quality/assembly.py new file mode 100644 index 00000000..3c49c1bb --- /dev/null +++ b/sciencebeam_parser/training/quality/assembly.py @@ -0,0 +1,199 @@ +"""Joining the record generation wrote with the counts only assembly can take. + +Generation sees the JATS and the TEI; the delft conversion sees the TEI and the +labels it produces, and is the only place the last count exists. It reads what +generation recorded rather than extending it, so that generated output stays +reproducible from generation alone and can be assembled more than once. +""" +import json +import logging +import os +from dataclasses import dataclass, field +from typing import Any, Dict, List, Mapping, Optional, Sequence + +from sciencebeam_parser.utils.io import auto_uploading_output_file, glob + + +LOGGER = logging.getLogger(__name__) + + +@dataclass +class GeneratedDocumentRecord: + """A row generation wrote, and which corpus it was read from.""" + document_id: str + corpus: Optional[str] + json_dict: Mapping[str, Any] + + @property + def jats_reference_count(self) -> Optional[int]: + return self.json_dict.get('jats', {}).get('reference_count') + + @property + def entity_element_count(self) -> Optional[int]: + return self.json_dict.get('entity_element_count') + + @property + def has_generated_output(self) -> bool: + """Whether generation wrote a file for this document at all. + + A row with no `written` at all is a document that failed or timed out. + """ + return self.json_dict.get('written') is True + + +def get_corpus_name_for_record_file_path(record_file_path: str) -> Optional[str]: + """The corpus a record belongs to, which is the directory above its model's.""" + model_directory = os.path.dirname(record_file_path) + corpus_directory = os.path.dirname(model_directory) + return os.path.basename(corpus_directory) or None + + +def read_generated_document_records( + record_path_pattern: str +) -> Dict[str, GeneratedDocumentRecord]: + record_file_list = glob(record_path_pattern) + if not record_file_list: + raise RuntimeError( + 'no quality record found for file pattern %r' % record_path_pattern + ) + LOGGER.info('reading quality records from: %r', record_file_list) + record_by_document_id: Dict[str, GeneratedDocumentRecord] = {} + for record_file_path in record_file_list: + corpus = get_corpus_name_for_record_file_path(record_file_path) + with open(record_file_path, 'r', encoding='utf-8') as record_file: + for line in record_file: + if not line.strip(): + continue + json_dict = json.loads(line) + record_by_document_id[json_dict['document_id']] = GeneratedDocumentRecord( + document_id=json_dict['document_id'], + corpus=corpus, + json_dict=json_dict, + ) + return record_by_document_id + + +@dataclass +class AssembledDocumentRecord: + """What assembly measured for one document, beside what generation recorded.""" + document_id: str + model_name: str + corpus: Optional[str] = None + sequence_count: int = 0 + entity_start_count: Optional[int] = None + label_start_counts: Optional[Dict[str, int]] = None + generated: Optional[GeneratedDocumentRecord] = None + + @property + def entity_element_count(self) -> Optional[int]: + return self.generated.entity_element_count if self.generated else None + + @property + def lost_at_parse(self) -> Optional[int]: + """Entities the parse did not return for an element the TEI holds.""" + element_count = self.entity_element_count + if element_count is None or self.entity_start_count is None: + return None + return element_count - self.entity_start_count + + def to_json_dict(self) -> Dict[str, Any]: + json_dict: Dict[str, Any] = { + 'document_id': self.document_id, + 'model': self.model_name, + 'corpus': self.corpus, + 'sequence_count': self.sequence_count, + } + if self.entity_start_count is not None: + json_dict['entity_start_count'] = self.entity_start_count + if self.label_start_counts: + json_dict['label_start_counts'] = self.label_start_counts + if self.generated is not None: + json_dict['generated'] = dict(self.generated.json_dict) + return json_dict + + +@dataclass +class CorpusAssemblySummary: + corpus: Optional[str] + document_count: int = 0 + sequence_count: int = 0 + entity_element_count: Optional[int] = None + entity_start_count: Optional[int] = None + documents_losing_entities: List[str] = field(default_factory=list) + documents_without_generated_record: List[str] = field(default_factory=list) + + def __str__(self) -> str: + parts = [ + f'{self.document_count} documents', + f'{self.sequence_count} sequences', + ] + # A model counted by label has no entity count, and reporting it as zero + # would read as every entity lost. + if self.entity_start_count is not None: + parts.append( + f'{self.entity_start_count} entities against ' + f'{self.entity_element_count or 0} elements' + ) + if self.documents_losing_entities: + parts.append( + f'{len(self.documents_losing_entities)} documents lost entities at ' + f'the parse: {sorted(self.documents_losing_entities)}' + ) + if self.documents_without_generated_record: + parts.append( + f'{len(self.documents_without_generated_record)} documents with no ' + f'record from generation' + ) + return '; '.join(parts) + + +def get_assembly_summary_by_corpus( + assembled_records: Sequence[AssembledDocumentRecord] +) -> Dict[Optional[str], CorpusAssemblySummary]: + summary_by_corpus: Dict[Optional[str], CorpusAssemblySummary] = {} + for record in assembled_records: + summary = summary_by_corpus.setdefault( + record.corpus, CorpusAssemblySummary(corpus=record.corpus) + ) + summary.document_count += 1 + summary.sequence_count += record.sequence_count + if record.entity_start_count is not None: + summary.entity_start_count = ( + (summary.entity_start_count or 0) + record.entity_start_count + ) + summary.entity_element_count = ( + (summary.entity_element_count or 0) + (record.entity_element_count or 0) + ) + if record.lost_at_parse: + summary.documents_losing_entities.append(record.document_id) + if record.generated is None: + summary.documents_without_generated_record.append(record.document_id) + return summary_by_corpus + + +def get_document_ids_without_generated_output( + record_by_document_id: Mapping[str, GeneratedDocumentRecord] +) -> List[str]: + """Documents generation recorded as having produced no file. + + Taken from the record's own `written`, not from what the training data is + missing: assembly is often pointed at part of a corpus, and inferring this + from the files present would report every document outside that part. + """ + return sorted( + document_id + for document_id, record in record_by_document_id.items() + if not record.has_generated_output + ) + + +def write_assembly_records( + output_file_path: str, + assembled_records: Sequence[AssembledDocumentRecord] +) -> None: + LOGGER.info('writing assembly quality record to: %r', output_file_path) + with auto_uploading_output_file( + output_file_path, mode='w', encoding='utf-8' + ) as output_file: + for record in assembled_records: + output_file.write(json.dumps(record.to_json_dict(), sort_keys=True) + '\n') diff --git a/sciencebeam_parser/training/quality/counting.py b/sciencebeam_parser/training/quality/counting.py index f74dbe15..17c5a8b5 100644 --- a/sciencebeam_parser/training/quality/counting.py +++ b/sciencebeam_parser/training/quality/counting.py @@ -5,11 +5,11 @@ references. """ import logging -from typing import Dict, FrozenSet, Iterable, Mapping, Optional, Sequence +from typing import Dict, FrozenSet, Iterable, Iterator, Mapping, Optional, Sequence from lxml import etree -from sciencebeam_parser.models.data import LayoutModelData +from sciencebeam_parser.models.data import LabeledLayoutToken, LayoutModelData from sciencebeam_parser.training.jats.field_vocab import CITATION_LABEL_BY_SUB_FIELD @@ -38,6 +38,32 @@ }) +# The label whose starts count one entity once the TEI is parsed back to labels, +# by model. This is the stage the training data ends up at, and the only one the +# delft conversion can see. +ENTITY_LABEL_BY_MODEL: Mapping[str, str] = { + 'reference-segmenter': '<reference>', +} + +# Models whose entities are one training sequence each, so that the entity count +# cannot change at the parse and what can is which labels are marked. +MODELS_COUNTED_BY_LABEL: FrozenSet[str] = frozenset({'citation'}) + + +def get_canonical_model_name(model_name: str) -> str: + """One spelling for a model, since the two CLIs disagree on it. + + `generate_data` names models as its generators do, hyphenated, and + `generate_delft_data` takes the underscored name the model registry uses. + A lookup that missed on the spelling would report no counts at all. + """ + return model_name.replace('_', '-') + + +def is_model_counted_by_label(model_name: str) -> bool: + return get_canonical_model_name(model_name) in MODELS_COUNTED_BY_LABEL + + def get_local_name(element: etree._Element) -> str: tag = element.tag if not isinstance(tag, str): @@ -47,9 +73,10 @@ def get_local_name(element: etree._Element) -> str: def count_entity_elements(model_name: str, tei_root: etree._Element) -> Optional[int]: """Return the number of entity elements the TEI holds, or None if the model has no count.""" - element_name = ENTITY_ELEMENT_NAME_BY_MODEL.get(model_name) + canonical_model_name = get_canonical_model_name(model_name) + element_name = ENTITY_ELEMENT_NAME_BY_MODEL.get(canonical_model_name) if element_name is None: - if model_name not in MODELS_WITHOUT_ENTITY_COUNT: + if canonical_model_name not in MODELS_WITHOUT_ENTITY_COUNT: LOGGER.warning( 'no cardinality check defined for model %r, and it is not declared as having none', model_name @@ -108,3 +135,50 @@ def _entry(label: str) -> Dict[str, int]: for label in get_labels_for_model_data_list(model_data_list): _entry(label)['marked'] += 1 return counts + + +def _iter_labels( + labeled_layout_tokens: Iterable[LabeledLayoutToken] +) -> Iterator[str]: + for labeled_layout_token in labeled_layout_tokens: + if labeled_layout_token.label: + yield labeled_layout_token.label + + +def count_entity_starts( + model_name: str, + labeled_layout_tokens_list: Sequence[Sequence[LabeledLayoutToken]] +) -> Optional[int]: + """Entities the training data ends up with, or None for a model counted by label. + + Fewer of these than the TEI holds elements is a boundary lost between + siblings; the elements are what generation recorded. + """ + entity_label = ENTITY_LABEL_BY_MODEL.get(get_canonical_model_name(model_name)) + if entity_label is None: + return None + return sum( + 1 + for labeled_layout_tokens in labeled_layout_tokens_list + for label in _iter_labels(labeled_layout_tokens) + if label == 'B-' + entity_label + ) + + +def count_label_starts_per_sequence( + labeled_layout_tokens_list: Sequence[Sequence[LabeledLayoutToken]] +) -> Dict[str, int]: + """Per label, the number of sequences that mark it. + + Presence per sequence, so it is comparable with what generation recorded as + marked for the citation model, where one sequence is one reference. + """ + counts: Dict[str, int] = {} + for labeled_layout_tokens in labeled_layout_tokens_list: + for label in { + label[2:] + for label in _iter_labels(labeled_layout_tokens) + if label.startswith('B-') + }: + counts[label] = counts.get(label, 0) + 1 + return counts diff --git a/tests/training/cli/generate_delft_data_test.py b/tests/training/cli/generate_delft_data_test.py index 2ba7d947..668157c5 100644 --- a/tests/training/cli/generate_delft_data_test.py +++ b/tests/training/cli/generate_delft_data_test.py @@ -1,4 +1,5 @@ # pylint: disable=not-callable +import json import logging import gzip from pathlib import Path @@ -27,6 +28,7 @@ import sciencebeam_parser.training.cli.generate_delft_data as generate_delft_data_module from sciencebeam_parser.training.cli.generate_delft_data import ( + get_document_id_for_tei_file, main, translate_tag_result_tags_IOB_to_grobid, translate_tags_IOB_to_grobid @@ -612,3 +614,122 @@ def test_should_be_able_to_load_and_generate_gzipped_training_data( LOGGER.debug('texts: %r', texts) assert len(texts) == 1 assert list(texts[0]) == tokens + + +@log_on_exception +class TestQualityRecord: + def _write_reference_segmenter_tei( + self, tei_source_path: Path, document_id: str, bibl_count: int + ) -> None: + tei_source_path.mkdir(parents=True, exist_ok=True) + ( + tei_source_path / f'{document_id}.references.referenceSegmenter.tei.xml' + ).write_bytes(etree.tostring(E('tei', E('text', E('listBibl', *[ + child + for index in range(bibl_count) + for child in (E('bibl', f'reference{index}', E('lb')), '\n') + ]))))) + + def test_should_record_the_entity_count_the_parse_returns(self, tmp_path: Path): + tei_source_path = tmp_path / 'tei' + self._write_reference_segmenter_tei(tei_source_path, 'document1', bibl_count=3) + output_path = tmp_path / 'output.data' + main([ + '--model-name=reference_segmenter', + f'--tei-source-path={tei_source_path}/*.tei.xml', + f'--delft-output-path={output_path}' + ]) + quality_record_path = Path(str(output_path) + '.quality.jsonl') + rows = [ + json.loads(line) + for line in quality_record_path.read_text(encoding='utf-8').splitlines() + ] + assert len(rows) == 1 + assert rows[0]['document_id'] == 'document1' + assert rows[0]['model'] == 'reference-segmenter' + assert rows[0]['entity_start_count'] == 3 + + def test_should_join_the_record_generation_wrote(self, tmp_path: Path): + tei_source_path = tmp_path / 'train' / 'ore' / 'reference-segmenter' / 'corpus' / 'tei' + self._write_reference_segmenter_tei(tei_source_path, 'document1', bibl_count=3) + generated_record_path = ( + tmp_path / 'train' / 'ore' / 'reference-segmenter' / 'quality.jsonl' + ) + generated_record_path.write_text(json.dumps({ + 'document_id': 'document1', + 'model': 'reference-segmenter', + 'status': 'ok', + 'jats': {'status': 'ok', 'reference_count': 4}, + 'written': True, + 'entity_element_count': 3, + }) + '\n', encoding='utf-8') + output_path = tmp_path / 'output.data' + main([ + '--model-name=reference_segmenter', + f'--tei-source-path={tei_source_path}/*.tei.xml', + f'--quality-record-path={generated_record_path}', + f'--delft-output-path={output_path}' + ]) + row = json.loads( + Path(str(output_path) + '.quality.jsonl').read_text(encoding='utf-8') + ) + assert row['corpus'] == 'ore' + assert row['entity_start_count'] == 3 + assert row['generated']['entity_element_count'] == 3 + assert row['generated']['jats']['reference_count'] == 4 + + def test_should_write_the_record_where_asked(self, tmp_path: Path): + tei_source_path = tmp_path / 'tei' + self._write_reference_segmenter_tei(tei_source_path, 'document1', bibl_count=1) + quality_output_path = tmp_path / 'elsewhere' / 'quality.jsonl' + main([ + '--model-name=reference_segmenter', + f'--tei-source-path={tei_source_path}/*.tei.xml', + f'--delft-output-path={tmp_path}/output.data', + f'--quality-output-path={quality_output_path}' + ]) + assert quality_output_path.exists() + + def test_should_record_the_labels_a_citation_sequence_marks(self, tmp_path: Path): + tei_source_path = tmp_path / 'tei' + tei_source_path.mkdir(parents=True) + (tei_source_path / 'document1.references.tei.xml').write_bytes(etree.tostring( + TEI_E('TEI', TEI_E('text', TEI_E('back', TEI_E('listBibl', *[ + TEI_E('bibl', TEI_E('title', TOKEN_1, {'level': 'a'}), ' ', TOKEN_2), + '\n', + ])))) + )) + output_path = tmp_path / 'output.data' + main([ + '--model-name=citation', + f'--tei-source-path={tei_source_path}/*.tei.xml', + f'--delft-output-path={output_path}' + ]) + row = json.loads( + Path(str(output_path) + '.quality.jsonl').read_text(encoding='utf-8') + ) + assert row['sequence_count'] == 1 + assert row['label_start_counts'] == {'<title>': 1} + # every bibl is its own sequence, so there is no entity count to take + assert 'entity_start_count' not in row + + +class TestGetDocumentIdForTeiFile: + def test_should_strip_the_model_suffix(self): + assert get_document_id_for_tei_file( + '/tei/PPR459453.references.referenceSegmenter.tei.xml', + '.references.referenceSegmenter.tei.xml' + ) == 'PPR459453' + + def test_should_strip_a_gzip_suffix_first(self): + assert get_document_id_for_tei_file( + '/tei/PPR459453.references.tei.xml.gz', '.references.tei.xml' + ) == 'PPR459453' + + def test_should_fall_back_for_a_model_with_no_declared_suffix(self): + assert get_document_id_for_tei_file('/tei/PPR459453.tei.xml', None) == 'PPR459453' + + def test_should_fall_back_when_the_file_does_not_carry_the_suffix(self): + assert get_document_id_for_tei_file( + '/tei/PPR459453.something-else.tei.xml', '.references.tei.xml' + ) == 'PPR459453' diff --git a/tests/training/quality/assembly_test.py b/tests/training/quality/assembly_test.py new file mode 100644 index 00000000..12ac02b1 --- /dev/null +++ b/tests/training/quality/assembly_test.py @@ -0,0 +1,237 @@ +import json +from pathlib import Path + +import pytest + +from sciencebeam_parser.training.quality.assembly import ( + AssembledDocumentRecord, + GeneratedDocumentRecord, + get_assembly_summary_by_corpus, + get_corpus_name_for_record_file_path, + get_document_ids_without_generated_output, + read_generated_document_records, + write_assembly_records +) + + +DOCUMENT_ID_1 = 'document1' +REFERENCE_SEGMENTER = 'reference-segmenter' + + +def _generated( + document_id: str = DOCUMENT_ID_1, + corpus: str = 'ore', + reference_count: int = 45, + entity_element_count: int = 45, + written: bool = True +) -> GeneratedDocumentRecord: + return GeneratedDocumentRecord( + document_id=document_id, + corpus=corpus, + json_dict={ + 'document_id': document_id, + 'model': REFERENCE_SEGMENTER, + 'jats': {'status': 'ok', 'reference_count': reference_count}, + 'written': written, + 'entity_element_count': entity_element_count, + }, + ) + + +class TestGetCorpusNameForRecordFilePath: + def test_should_take_the_directory_above_the_model(self): + assert get_corpus_name_for_record_file_path( + '/data/train/ore/reference-segmenter/quality.jsonl' + ) == 'ore' + + def test_should_be_none_without_a_directory_to_read(self): + assert get_corpus_name_for_record_file_path('quality.jsonl') is None + + +class TestReadGeneratedDocumentRecords: + def test_should_read_the_rows_and_the_corpus_they_came_from(self, tmp_path: Path): + record_file_path = tmp_path / 'train' / 'ore' / REFERENCE_SEGMENTER / 'quality.jsonl' + record_file_path.parent.mkdir(parents=True) + record_file_path.write_text('\n'.join([ + json.dumps({'document_id': 'document1', 'entity_element_count': 12}), + json.dumps({'document_id': 'document2', 'entity_element_count': 34}), + ]) + '\n', encoding='utf-8') + record_by_document_id = read_generated_document_records(str(record_file_path)) + assert set(record_by_document_id) == {'document1', 'document2'} + assert record_by_document_id['document1'].corpus == 'ore' + assert record_by_document_id['document2'].entity_element_count == 34 + + def test_should_fail_when_no_record_matches(self, tmp_path: Path): + with pytest.raises(RuntimeError): + read_generated_document_records(str(tmp_path / 'not-there' / '*.jsonl')) + + +class TestAssembledDocumentRecord: + def test_should_report_entities_lost_at_the_parse(self): + record = AssembledDocumentRecord( + document_id=DOCUMENT_ID_1, + model_name=REFERENCE_SEGMENTER, + entity_start_count=1, + generated=_generated(entity_element_count=45), + ) + assert record.lost_at_parse == 44 + + def test_should_report_nothing_lost_when_every_element_returned_an_entity(self): + record = AssembledDocumentRecord( + document_id=DOCUMENT_ID_1, + model_name=REFERENCE_SEGMENTER, + entity_start_count=45, + generated=_generated(entity_element_count=45), + ) + assert record.lost_at_parse == 0 + + def test_should_not_claim_a_loss_without_a_record_from_generation(self): + record = AssembledDocumentRecord( + document_id=DOCUMENT_ID_1, + model_name=REFERENCE_SEGMENTER, + entity_start_count=45, + ) + assert record.lost_at_parse is None + + def test_should_carry_the_generated_counts_into_its_row(self): + json_dict = AssembledDocumentRecord( + document_id=DOCUMENT_ID_1, + model_name=REFERENCE_SEGMENTER, + corpus='ore', + sequence_count=1, + entity_start_count=45, + generated=_generated(), + ).to_json_dict() + assert json_dict['corpus'] == 'ore' + assert json_dict['entity_start_count'] == 45 + assert json_dict['generated']['jats']['reference_count'] == 45 + + +class TestGetAssemblySummaryByCorpus: + def test_should_total_the_counts_per_corpus(self): + summary_by_corpus = get_assembly_summary_by_corpus([ + AssembledDocumentRecord( + document_id='document1', model_name=REFERENCE_SEGMENTER, corpus='ore', + sequence_count=1, entity_start_count=10, + generated=_generated('document1', entity_element_count=10), + ), + AssembledDocumentRecord( + document_id='document2', model_name=REFERENCE_SEGMENTER, corpus='ore', + sequence_count=1, entity_start_count=20, + generated=_generated('document2', entity_element_count=20), + ), + ]) + summary = summary_by_corpus['ore'] + assert summary.document_count == 2 + assert summary.entity_start_count == 30 + assert summary.entity_element_count == 30 + assert not summary.documents_losing_entities + + def test_should_name_the_documents_that_lost_entities(self): + summary_by_corpus = get_assembly_summary_by_corpus([ + AssembledDocumentRecord( + document_id='collapsed', model_name=REFERENCE_SEGMENTER, corpus='ore', + sequence_count=1, entity_start_count=1, + generated=_generated('collapsed', entity_element_count=40), + ), + AssembledDocumentRecord( + document_id='intact', model_name=REFERENCE_SEGMENTER, corpus='ore', + sequence_count=1, entity_start_count=40, + generated=_generated('intact', entity_element_count=40), + ), + ]) + summary = summary_by_corpus['ore'] + assert summary.documents_losing_entities == ['collapsed'] + assert 'collapsed' in str(summary) + + def test_should_count_a_document_with_no_record_from_generation(self): + summary_by_corpus = get_assembly_summary_by_corpus([ + AssembledDocumentRecord( + document_id='unrecorded', model_name=REFERENCE_SEGMENTER, + sequence_count=1, entity_start_count=3, + ), + ]) + assert summary_by_corpus[None].documents_without_generated_record == ['unrecorded'] + + def test_should_keep_corpora_apart(self): + summary_by_corpus = get_assembly_summary_by_corpus([ + AssembledDocumentRecord( + document_id='document1', model_name=REFERENCE_SEGMENTER, corpus='ore', + sequence_count=1, entity_start_count=10, + ), + AssembledDocumentRecord( + document_id='document2', model_name=REFERENCE_SEGMENTER, + corpus='scielo_preprints-jats', + sequence_count=1, entity_start_count=20, + ), + ]) + assert summary_by_corpus['ore'].entity_start_count == 10 + assert summary_by_corpus['scielo_preprints-jats'].entity_start_count == 20 + + +class TestGetDocumentIdsWithoutGeneratedOutput: + def test_should_find_a_document_generation_wrote_no_file_for(self): + assert get_document_ids_without_generated_output({ + 'document1': _generated('document1'), + 'document2': _generated('document2', written=False), + }) == ['document2'] + + def test_should_find_a_document_that_failed_before_writing_anything(self): + record = GeneratedDocumentRecord( + document_id='timed-out', + corpus='ore', + json_dict={'document_id': 'timed-out', 'status': 'timeout'}, + ) + assert get_document_ids_without_generated_output({ + 'timed-out': record + }) == ['timed-out'] + + def test_should_find_nothing_when_every_document_was_written(self): + assert not get_document_ids_without_generated_output({ + 'document1': _generated('document1') + }) + + def test_should_not_depend_on_which_documents_were_assembled(self): + # Assembly is often pointed at part of a corpus; the record is what says + # whether a file was written, so a narrowed run must not report the rest. + assert not get_document_ids_without_generated_output({ + 'document1': _generated('document1'), + 'document2': _generated('document2'), + }) + + +class TestWriteAssemblyRecords: + def test_should_write_one_json_line_per_document(self, tmp_path: Path): + output_file_path = tmp_path / 'reference-segmenter.data.quality.jsonl' + write_assembly_records(str(output_file_path), [ + AssembledDocumentRecord( + document_id='document1', model_name=REFERENCE_SEGMENTER, + sequence_count=1, entity_start_count=10, + ), + AssembledDocumentRecord( + document_id='document2', model_name=REFERENCE_SEGMENTER, + sequence_count=1, entity_start_count=20, + ), + ]) + lines = output_file_path.read_text(encoding='utf-8').splitlines() + assert [json.loads(line)['document_id'] for line in lines] == [ + 'document1', 'document2' + ] + assert [json.loads(line)['entity_start_count'] for line in lines] == [10, 20] + + +class TestCorpusAssemblySummaryForLabelCountedModel: + def test_should_not_report_a_missing_entity_count_as_zero_entities(self): + # Every citation element is its own sequence, so it has no entity count; + # reporting it as zero would read as every entity lost. + summary = get_assembly_summary_by_corpus([ + AssembledDocumentRecord( + document_id=DOCUMENT_ID_1, model_name='citation', corpus='ore', + sequence_count=45, entity_start_count=None, + label_start_counts={'<title>': 44}, + generated=_generated(entity_element_count=45), + ), + ])['ore'] + assert summary.entity_start_count is None + assert 'entities' not in str(summary) + assert '45 sequences' in str(summary) diff --git a/tests/training/quality/counting_test.py b/tests/training/quality/counting_test.py index 564045ef..d8c9c199 100644 --- a/tests/training/quality/counting_test.py +++ b/tests/training/quality/counting_test.py @@ -1,12 +1,20 @@ from lxml import etree from sciencebeam_parser.models.citation.labels import IDENTIFIER_LABEL -from sciencebeam_parser.models.data import LabeledLayoutModelData, LayoutModelData +from sciencebeam_parser.document.layout_document import LayoutToken +from sciencebeam_parser.models.data import ( + LabeledLayoutModelData, + LabeledLayoutToken, + LayoutModelData +) from sciencebeam_parser.training.jats.field_vocab import JatsSubFieldNames from sciencebeam_parser.training.quality.counting import ( count_citation_labels, count_entity_elements, - get_labels_for_model_data_list + count_entity_starts, + count_label_starts_per_sequence, + get_labels_for_model_data_list, + is_model_counted_by_label ) @@ -120,3 +128,79 @@ def test_should_count_over_all_references(self): ) assert counts['<title>'] == {'jats': 2, 'marked': 1} assert counts['<date>'] == {'jats': 1, 'marked': 2} + + +def _labeled_token(label: str) -> LabeledLayoutToken: + return LabeledLayoutToken(label=label, layout_token=LayoutToken('token')) + + +class TestCountEntityStarts: + def test_should_count_one_entity_per_start(self): + assert count_entity_starts('reference-segmenter', [[ + _labeled_token('B-<reference>'), + _labeled_token('I-<reference>'), + _labeled_token('B-<reference>'), + ]]) == 2 + + def test_should_count_across_sequences(self): + assert count_entity_starts('reference-segmenter', [ + [_labeled_token('B-<reference>')], + [_labeled_token('B-<reference>')], + ]) == 2 + + def test_should_not_count_another_label(self): + assert count_entity_starts('reference-segmenter', [[ + _labeled_token('B-<label>'), + _labeled_token('B-<reference>'), + ]]) == 1 + + def test_should_count_an_element_that_returned_no_entity_as_missing(self): + # Two elements sharing a label came back as one entity before the parse + # kept element boundaries; the count is what says so. + assert count_entity_starts('reference-segmenter', [[ + _labeled_token('B-<reference>'), + _labeled_token('I-<reference>'), + _labeled_token('I-<reference>'), + ]]) == 1 + + def test_should_return_none_for_a_model_counted_by_label(self): + assert count_entity_starts('citation', [[_labeled_token('B-<title>')]]) is None + + +class TestCountLabelStartsPerSequence: + def test_should_count_a_label_once_per_sequence(self): + assert count_label_starts_per_sequence([[ + _labeled_token('B-<author>'), + _labeled_token('I-<author>'), + _labeled_token('B-<author>'), + ]]) == {'<author>': 1} + + def test_should_count_each_sequence_that_marks_the_label(self): + assert count_label_starts_per_sequence([ + [_labeled_token('B-<title>'), _labeled_token('B-<date>')], + [_labeled_token('B-<title>')], + ]) == {'<title>': 2, '<date>': 1} + + def test_should_ignore_tokens_outside_any_entity(self): + assert count_label_starts_per_sequence([[ + _labeled_token('O'), _labeled_token('B-<title>'), _labeled_token('O'), + ]]) == {'<title>': 1} + + +class TestGetCanonicalModelName: + def test_should_accept_the_underscored_spelling_the_delft_cli_takes(self): + # generate_delft_data is invoked with --model-name=reference_segmenter, + # while generation names the same model reference-segmenter. + assert count_entity_starts('reference_segmenter', [[ + _labeled_token('B-<reference>'), _labeled_token('B-<reference>'), + ]]) == 2 + + def test_should_accept_the_underscored_spelling_for_elements(self): + tei_root = etree.fromstring( + '<tei><text><listBibl><bibl>Reference 1</bibl></listBibl></text></tei>' + ) + assert count_entity_elements('reference_segmenter', tei_root) == 1 + + def test_should_recognise_a_label_counted_model_either_way(self): + assert is_model_counted_by_label('citation') + assert not is_model_counted_by_label('reference_segmenter') From deb1ee79d96b8826986decf411f544689ebdfc2b Mon Sep 17 00:00:00 2001 From: Daniel Ecer <d.ecer@elifesciences.org> Date: Thu, 20 Aug 2026 13:57:01 +0100 Subject: [PATCH 3/4] Exclude training data that fails a stated threshold The record measured; nothing acted on it. A corpus half of whose documents present their reference list as one reference could still reach a training run, which is what happened. generate_delft_data now takes --quality-filter, leaving out the documents that fail the thresholds in resources/training_quality.yml and reporting each with the stage that failed and the numbers behind it, in the log and in the row the training run can read. Assembly refuses rather than proceeds when a corpus loses more than a configured share of its documents: --max-excluded-ratio overrides that for a run, which is a decision to record rather than a way around it. Thresholds are per model, and a model with no entry fails rather than being assumed sound; the models whose labels mark regions rather than repeated entities carry cardinality: none and a reason. Two things are recorded rather than failed, both measured rather than assumed. More elements than the JATS has references is the reference segmenter writing a block per contiguous run of a reference's lines, so a reference split across a column becomes two elements where the citation model's count of the same references matches the JATS exactly; those seven scielo_preprints-jats documents are otherwise sound and the split belongs to alignment. And no citation label carries a floor, because the level of a label's rate says as much about the publisher as the pipeline: ORE prints no DOI at all where its JATS carries one for 1205 of 1679 references, so a floor on the identifier would reject nearly every one of its documents. A rate that moves is the finding. Filtering is off unless asked for, since this command is also run over corpora that carry no record at all, and what cannot be checked is reported as unchecked rather than assumed good. Over the committed corpora it keeps 37 of 38 ore documents and 49 of 50 scielo_preprints-jats, excluding 5-264_v2 for a JATS that will not parse -- 29 unlabelled citation references, since region and reference boundaries come from the models while labels come from the JATS -- and PPR459453 for a reference region holding 2 elements against 45 references. The reference segmenter's data goes from 88 documents to 86 and citation's from 3335 sequences to 3304. log_on_exception dropped the return value of whatever it wrapped, which is invisible for a test method and silently returns None for a helper in a decorated class. --- doc/training.md | 42 +++ .../resources/training_quality.yml | 86 ++++++ .../training/cli/generate_delft_data.py | 127 ++++++++- .../training/quality/assembly.py | 15 ++ sciencebeam_parser/training/quality/gate.py | 247 ++++++++++++++++++ tests/test_utils.py | 3 +- .../training/cli/generate_delft_data_test.py | 125 +++++++++ tests/training/quality/gate_test.py | 223 ++++++++++++++++ 8 files changed, 856 insertions(+), 12 deletions(-) create mode 100644 sciencebeam_parser/resources/training_quality.yml create mode 100644 sciencebeam_parser/training/quality/gate.py create mode 100644 tests/training/quality/gate_test.py diff --git a/doc/training.md b/doc/training.md index 24e9e7b3..913a9d95 100644 --- a/doc/training.md +++ b/doc/training.md @@ -185,6 +185,48 @@ python -m sciencebeam_parser.training.cli.generate_delft_data \ Without `--quality-record-path` the counts are still recorded, with nothing to compare them against — which is enough to re-check a committed corpus offline. +#### Filtering on quality + +`--quality-filter` leaves out the documents that fail the thresholds in +[`training_quality.yml`](../sciencebeam_parser/resources/training_quality.yml), +rather than only recording their counts. Each exclusion is reported with the stage +that failed and the numbers behind it, and the summary says what was kept per +corpus: + +```text +excluding 5-264_v2: excluded (jats-not-readable) [jats_status=unparsable] +excluding PPR459453: excluded (elements-short-of-jats) + [element_ratio=0.044, entity_element_count=2, jats_reference_count=45] +ore / reference-segmenter: kept 37 of 38 documents (3% excluded); + jats-not-readable: ['5-264_v2'] +``` + +The same reasons are written to each row of the record, so a corpus that shrinks +can always be accounted for document by document. + +Assembly **refuses** rather than proceeds when a corpus loses more than +`corpus.max_excluded_ratio` of its documents: dropping most of a corpus is a +finding about the pipeline, not a routine filter outcome. `--max-excluded-ratio` +overrides it for a run, which is a decision to record rather than a way around +the refusal. + +Thresholds are per model, and a model with no entry in the config fails rather +than being assumed sound. Models whose labels mark regions instead of repeated +entities carry `cardinality: none` and a reason. Two things are deliberately +recorded rather than failed: + +- **more elements than the JATS has references.** The reference segmenter writes a + block per contiguous run of a reference's lines, so a reference split across a + column becomes two elements — where the citation model's count of the same + references matches the JATS exactly. The document is otherwise sound. +- **citation label rates.** No floor is set, because the level of a label's rate + says as much about the publisher as the pipeline: ORE prints no DOI at all where + its JATS carries one for 1205 of 1679 references, so a floor on the identifier + would reject nearly every ORE document. A rate that moves is the finding. + +Filtering is off unless asked for, since this command is also run over corpora +that carry no record at all, such as GROBID's own. + #### Example command for `segmentation` model ```bash diff --git a/sciencebeam_parser/resources/training_quality.yml b/sciencebeam_parser/resources/training_quality.yml new file mode 100644 index 00000000..2a89b0dd --- /dev/null +++ b/sciencebeam_parser/resources/training_quality.yml @@ -0,0 +1,86 @@ +# What the quality gate requires of a generated document before its training +# data is used, per model. +# +# The counts come from the quality record: the references the JATS declares, the +# elements generation wrote, and the entities those elements parse back to. Each +# threshold is a floor on one stage against the one before it, so a document that +# fails names the stage that lost its references rather than only that it is wrong. +# +# min_jats_reference_count a document whose JATS declares fewer references than +# this has no reference of record; no comparison +# between stages can flag it, because every stage +# agrees on zero +# min_element_ratio elements written, over the references the JATS +# declares. An exact match is too strict: what a +# PDF's reference list contains and what the JATS +# carries legitimately differ +# min_entity_ratio entities the training data ends up with, over the +# elements the TEI holds. This is the stage where two +# sibling elements sharing a label used to come back as +# one entity, and the floor is what keeps that from +# returning unnoticed +# +# Over-counting is recorded and not failed. The reference segmenter writes a +# block per contiguous run of a reference's lines, so a reference split across a +# column or a page becomes two elements -- 7 of 50 scielo_preprints-jats documents, +# where the citation model's count of the same references matches the JATS exactly. +# Those documents are otherwise sound, and the split is a defect of alignment +# rather than of the data's cardinality. +# +# A model with no cardinality to check carries `cardinality: none` and a reason, +# rather than being left out: a missing entry is a decision not taken, and the +# gate says so instead of passing silently. + +corpus: + # Beyond this share of a corpus excluded, assembly refuses rather than + # proceeds. Dropping most of a corpus is a finding about the pipeline, not a + # routine filter outcome, and the counterpart of failing on inconsistent + # feature lengths on the training side. Raise it deliberately, with the + # reason recorded, rather than working around the refusal. + max_excluded_ratio: 0.2 + +models: + reference-segmenter: + min_jats_reference_count: 1 + min_element_ratio: 0.9 + min_entity_ratio: 0.8 + + citation: + min_jats_reference_count: 1 + min_element_ratio: 0.9 + # Every element is its own training sequence, so there is no entity ratio to + # take; what can change is which labels are marked, and that is recorded per + # label rather than gated. No floor is set because the level of a label's + # rate says as much about the publisher as the pipeline: title, journal, + # date, volume, pages and author sit at 0.95 and above on both measured + # corpora, while ORE prints no DOI at all where its JATS carries one for + # 1205 of 1679 references, so a floor on the identifier would reject nearly + # every one of its documents. A rate that moves is the finding. + label_floors: {} + + segmentation: + cardinality: none + reason: labels mark regions of a document, which occur once rather than repeatedly + header: + cardinality: none + reason: labels mark fields of one header, so presence is what matters + affiliation-address: + cardinality: none + reason: > + one entity per affiliation, but the JATS count is of affiliations rather than + references and has no measured threshold yet + name-header: + cardinality: none + reason: one entity per author name, with no measured threshold yet + name-citation: + cardinality: none + reason: one entity per author name, with no measured threshold yet + fulltext: + cardinality: none + reason: labels mark body regions, which occur once each + figure: + cardinality: none + reason: one entity per figure, with no measured threshold yet + table: + cardinality: none + reason: one entity per table, with no measured threshold yet diff --git a/sciencebeam_parser/training/cli/generate_delft_data.py b/sciencebeam_parser/training/cli/generate_delft_data.py index 84defe58..eb0d0060 100644 --- a/sciencebeam_parser/training/cli/generate_delft_data.py +++ b/sciencebeam_parser/training/cli/generate_delft_data.py @@ -47,6 +47,14 @@ read_generated_document_records, write_assembly_records ) +from sciencebeam_parser.training.quality.gate import ( + TRAINING_QUALITY_CONFIG_FILE, + TrainingQualityConfig, + check_corpus_loss_or_fail, + get_gate_summary_by_corpus, + get_quality_verdict, + load_training_quality_config +) from sciencebeam_parser.training.quality.counting import ( count_entity_starts, count_label_starts_per_sequence, @@ -105,6 +113,30 @@ def parse_args(argv: Optional[List[str]] = None) -> argparse.Namespace: ' (default: the delft output path with ".quality.jsonl" appended).' ) ) + parser.add_argument( + '--quality-filter', + action='store_true', + help=( + 'Exclude documents failing the quality thresholds, rather than only' + ' recording their counts. Assembly refuses when a corpus loses more than' + ' the configured share of its documents.' + ) + ) + parser.add_argument( + '--quality-config-path', + type=str, + default=TRAINING_QUALITY_CONFIG_FILE, + help='Thresholds the quality filter applies (default: the shipped config).' + ) + parser.add_argument( + '--max-excluded-ratio', + type=float, + required=False, + help=( + 'Override the configured share of a corpus beyond which assembly refuses.' + ' Raising it is a decision to record, not a way around the refusal.' + ) + ) parser.add_argument( '--include-extra-columns', action='store_true', @@ -346,22 +378,68 @@ def get_assembled_document_record( model_name: str, result: DelftDocumentResult, generated_record_by_document_id: Mapping[str, GeneratedDocumentRecord], + quality_config: Optional[TrainingQualityConfig] = None, ) -> AssembledDocumentRecord: generated = generated_record_by_document_id.get(document_id) + entity_start_count = count_entity_starts( + model_name, result.labeled_layout_tokens_list + ) + verdict = None + if quality_config is not None: + verdict = get_quality_verdict( + document_id=document_id, + thresholds=quality_config.get_thresholds_for_model( + get_canonical_model_name(model_name) + ), + jats_status=generated.jats_status if generated else None, + jats_reference_count=generated.jats_reference_count if generated else None, + written=generated.written if generated else None, + entity_element_count=generated.entity_element_count if generated else None, + entity_start_count=entity_start_count, + sequence_count=len(result.labeled_layout_tokens_list), + ) return AssembledDocumentRecord( document_id=document_id, model_name=get_canonical_model_name(model_name), corpus=generated.corpus if generated else None, sequence_count=len(result.labeled_layout_tokens_list), - entity_start_count=count_entity_starts( - model_name, result.labeled_layout_tokens_list - ), + entity_start_count=entity_start_count, label_start_counts=( count_label_starts_per_sequence(result.labeled_layout_tokens_list) if is_model_counted_by_label(model_name) else None ), generated=generated, + verdict=verdict, + ) + + +def log_gate_summary( + model_name: str, + assembled_records: Sequence[AssembledDocumentRecord], + quality_config: TrainingQualityConfig, + max_excluded_ratio: Optional[float] = None, +) -> None: + summary_by_corpus = get_gate_summary_by_corpus([ + (record.verdict, record.corpus) + for record in assembled_records + if record.verdict is not None + ]) + for record in assembled_records: + if record.verdict is not None and record.verdict.is_excluded: + LOGGER.warning('excluding %s', record.verdict) + for corpus, summary in sorted( + summary_by_corpus.items(), key=lambda item: item[0] or '' + ): + LOGGER.info( + '%s / %s: %s', corpus or 'corpus not known', + get_canonical_model_name(model_name), summary + ) + check_corpus_loss_or_fail( + summary_by_corpus, + max_excluded_ratio + if max_excluded_ratio is not None + else quality_config.max_excluded_ratio ) @@ -397,7 +475,10 @@ def generate_delft_training_data( # pylint: disable=too-many-locals sciencebeam_parser: ScienceBeamParser, include_extra_columns: bool = False, quality_record_path: Optional[str] = None, - quality_output_path: Optional[str] = None + quality_output_path: Optional[str] = None, + quality_filter: bool = False, + quality_config_path: str = TRAINING_QUALITY_CONFIG_FILE, + max_excluded_ratio: Optional[float] = None ): training_tei_parser = get_training_tei_parser_for_model_name( model_name, @@ -434,6 +515,14 @@ def generate_delft_training_data( # pylint: disable=too-many-locals tei_filename_suffix = get_tei_filename_suffix_for_model_name( model_name, sciencebeam_parser=sciencebeam_parser ) + quality_config = ( + load_training_quality_config(quality_config_path) if quality_filter else None + ) + if quality_filter and not quality_record_path: + LOGGER.warning( + 'no --quality-record-path given, so the thresholds against the JATS and' + ' the generated elements cannot be applied and are reported as unchecked' + ) assembled_records: List[AssembledDocumentRecord] = [] LOGGER.info('writing to : %r', delft_output_path) with auto_uploading_output_file( @@ -441,9 +530,8 @@ def generate_delft_training_data( # pylint: disable=too-many-locals mode='w', encoding='utf-8', ) as data_fp: - for document_index, (tei_file, raw_file) in enumerate(zip(tei_file_list, raw_file_list)): - if document_index > 0: - data_fp.write('\n\n') + written_document_count = 0 + for tei_file, raw_file in zip(tei_file_list, raw_file_list): result = get_delft_training_data_for_document( tei_file=tei_file, raw_file=raw_file, @@ -452,13 +540,22 @@ def generate_delft_training_data( # pylint: disable=too-many-locals column_layout=column_layout, include_extra_columns=include_extra_columns ) - data_fp.writelines(result.data_lines) - assembled_records.append(get_assembled_document_record( + record = get_assembled_document_record( document_id=get_document_id_for_tei_file(tei_file, tei_filename_suffix), model_name=model_name, result=result, generated_record_by_document_id=generated_record_by_document_id, - )) + quality_config=quality_config, + ) + assembled_records.append(record) + if record.verdict is not None and record.verdict.is_excluded: + continue + # The blank line separates documents, so it follows what was written + # rather than the position in the source list. + if written_document_count: + data_fp.write('\n\n') + data_fp.writelines(result.data_lines) + written_document_count += 1 write_assembly_records( quality_output_path or delft_output_path + '.quality.jsonl', assembled_records @@ -466,6 +563,11 @@ def generate_delft_training_data( # pylint: disable=too-many-locals log_assembly_summary( model_name, assembled_records, generated_record_by_document_id ) + if quality_config is not None: + log_gate_summary( + model_name, assembled_records, quality_config, + max_excluded_ratio=max_excluded_ratio + ) def run(args: argparse.Namespace): @@ -482,7 +584,10 @@ def run(args: argparse.Namespace): sciencebeam_parser=sciencebeam_parser, include_extra_columns=args.include_extra_columns, quality_record_path=args.quality_record_path, - quality_output_path=args.quality_output_path + quality_output_path=args.quality_output_path, + quality_filter=args.quality_filter, + quality_config_path=args.quality_config_path, + max_excluded_ratio=args.max_excluded_ratio ) diff --git a/sciencebeam_parser/training/quality/assembly.py b/sciencebeam_parser/training/quality/assembly.py index 3c49c1bb..672d5149 100644 --- a/sciencebeam_parser/training/quality/assembly.py +++ b/sciencebeam_parser/training/quality/assembly.py @@ -11,6 +11,7 @@ from dataclasses import dataclass, field from typing import Any, Dict, List, Mapping, Optional, Sequence +from sciencebeam_parser.training.quality.gate import QualityVerdict from sciencebeam_parser.utils.io import auto_uploading_output_file, glob @@ -32,6 +33,14 @@ def jats_reference_count(self) -> Optional[int]: def entity_element_count(self) -> Optional[int]: return self.json_dict.get('entity_element_count') + @property + def written(self) -> Optional[bool]: + return self.json_dict.get('written') + + @property + def jats_status(self) -> Optional[str]: + return self.json_dict.get('jats', {}).get('status') + @property def has_generated_output(self) -> bool: """Whether generation wrote a file for this document at all. @@ -83,6 +92,7 @@ class AssembledDocumentRecord: entity_start_count: Optional[int] = None label_start_counts: Optional[Dict[str, int]] = None generated: Optional[GeneratedDocumentRecord] = None + verdict: Optional['QualityVerdict'] = None @property def entity_element_count(self) -> Optional[int]: @@ -109,6 +119,11 @@ def to_json_dict(self) -> Dict[str, Any]: json_dict['label_start_counts'] = self.label_start_counts if self.generated is not None: json_dict['generated'] = dict(self.generated.json_dict) + if self.verdict is not None: + json_dict['excluded'] = self.verdict.is_excluded + if self.verdict.is_excluded: + json_dict['exclusion_reasons'] = list(self.verdict.exclusion_reasons) + json_dict['exclusion_detail'] = dict(self.verdict.detail) return json_dict diff --git a/sciencebeam_parser/training/quality/gate.py b/sciencebeam_parser/training/quality/gate.py new file mode 100644 index 00000000..523deda0 --- /dev/null +++ b/sciencebeam_parser/training/quality/gate.py @@ -0,0 +1,247 @@ +"""Deciding what a training run may use, from the record and a stated threshold. + +The gate judges; the record measures. It excludes at assembly rather than at +generation, so the documents it refuses stay generatable and alignment stays +investigable. +""" +import logging +import os +from dataclasses import dataclass, field +from typing import Any, Dict, List, Mapping, Optional, Sequence + +import yaml + + +LOGGER = logging.getLogger(__name__) + + +TRAINING_QUALITY_CONFIG_FILE = os.path.join( + os.path.dirname( # sciencebeam_parser + os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + ), + 'resources', + 'training_quality.yml' +) + + +NO_CARDINALITY = 'none' + + +class ExclusionReason: + """Why a document's training data is not used, in the order the stages run.""" + JATS_NOT_READABLE = 'jats-not-readable' + JATS_HAS_NO_REFERENCES = 'jats-has-no-references' + NO_GENERATED_OUTPUT = 'no-generated-output' + ELEMENTS_SHORT_OF_JATS = 'elements-short-of-jats' + ENTITIES_SHORT_OF_ELEMENTS = 'entities-short-of-elements' + NO_TRAINING_SEQUENCES = 'no-training-sequences' + + +@dataclass(frozen=True) +class ModelQualityThresholds: + model_name: str + cardinality: Optional[str] = None + reason: Optional[str] = None + min_jats_reference_count: Optional[int] = None + min_element_ratio: Optional[float] = None + min_entity_ratio: Optional[float] = None + label_floors: Mapping[str, float] = field(default_factory=dict) + + @property + def has_cardinality_check(self) -> bool: + return self.cardinality != NO_CARDINALITY + + +@dataclass(frozen=True) +class TrainingQualityConfig: + max_excluded_ratio: float + thresholds_by_model: Mapping[str, ModelQualityThresholds] + + def get_thresholds_for_model(self, model_name: str) -> ModelQualityThresholds: + thresholds = self.thresholds_by_model.get(model_name) + if thresholds is None: + raise KeyError( + 'no quality thresholds configured for model %r; add an entry,' + ' with cardinality: %s and a reason if it has no count to check' + % (model_name, NO_CARDINALITY) + ) + return thresholds + + +def load_training_quality_config( + config_file_path: str = TRAINING_QUALITY_CONFIG_FILE +) -> TrainingQualityConfig: + with open(config_file_path, 'r', encoding='utf-8') as config_file: + config_json = yaml.safe_load(config_file) + return TrainingQualityConfig( + max_excluded_ratio=config_json['corpus']['max_excluded_ratio'], + thresholds_by_model={ + model_name: ModelQualityThresholds(model_name=model_name, **(entry or {})) + for model_name, entry in config_json['models'].items() + }, + ) + + +@dataclass +class QualityVerdict: + """Whether a document's training data is used, and the numbers behind it.""" + document_id: str + exclusion_reasons: Sequence[str] = () + detail: Mapping[str, Any] = field(default_factory=dict) + + @property + def is_excluded(self) -> bool: + return bool(self.exclusion_reasons) + + @property + def primary_reason(self) -> Optional[str]: + """The earliest stage that failed, which is the one to fix.""" + return self.exclusion_reasons[0] if self.exclusion_reasons else None + + def __str__(self) -> str: + if not self.is_excluded: + return f'{self.document_id}: kept' + detail = ', '.join(f'{key}={value}' for key, value in sorted(self.detail.items())) + return ( + f'{self.document_id}: excluded ({", ".join(self.exclusion_reasons)})' + f'{" [" + detail + "]" if detail else ""}' + ) + + +def _get_ratio(numerator: Optional[int], denominator: Optional[int]) -> Optional[float]: + if numerator is None or not denominator: + return None + return numerator / denominator + + +def get_quality_verdict( # pylint: disable=too-many-branches + document_id: str, + thresholds: ModelQualityThresholds, + jats_status: Optional[str] = None, + jats_reference_count: Optional[int] = None, + written: Optional[bool] = None, + entity_element_count: Optional[int] = None, + entity_start_count: Optional[int] = None, + sequence_count: int = 0, +) -> QualityVerdict: + """Judge one document, naming every stage that failed, earliest first. + + A count that is not available is not a failure: assembly is run over corpora + with no record at all, and what cannot be checked has to be reported as + unchecked rather than assumed good. + """ + reasons: List[str] = [] + detail: Dict[str, Any] = {} + if sequence_count == 0: + reasons.append(ExclusionReason.NO_TRAINING_SEQUENCES) + if not thresholds.has_cardinality_check: + return QualityVerdict(document_id, reasons, detail) + + if jats_status is not None and jats_status != 'ok': + reasons.insert(0, ExclusionReason.JATS_NOT_READABLE) + detail['jats_status'] = jats_status + elif ( + thresholds.min_jats_reference_count is not None + and jats_reference_count is not None + and jats_reference_count < thresholds.min_jats_reference_count + ): + reasons.insert(0, ExclusionReason.JATS_HAS_NO_REFERENCES) + detail['jats_reference_count'] = jats_reference_count + if written is False: + reasons.append(ExclusionReason.NO_GENERATED_OUTPUT) + + element_ratio = _get_ratio(entity_element_count, jats_reference_count) + if ( + thresholds.min_element_ratio is not None + and element_ratio is not None + and element_ratio < thresholds.min_element_ratio + ): + reasons.append(ExclusionReason.ELEMENTS_SHORT_OF_JATS) + detail['element_ratio'] = round(element_ratio, 3) + detail['entity_element_count'] = entity_element_count + detail['jats_reference_count'] = jats_reference_count + + entity_ratio = _get_ratio(entity_start_count, entity_element_count) + if ( + thresholds.min_entity_ratio is not None + and entity_ratio is not None + and entity_ratio < thresholds.min_entity_ratio + ): + reasons.append(ExclusionReason.ENTITIES_SHORT_OF_ELEMENTS) + detail['entity_ratio'] = round(entity_ratio, 3) + detail['entity_start_count'] = entity_start_count + detail['entity_element_count'] = entity_element_count + + return QualityVerdict(document_id, reasons, detail) + + +class CorpusMostlyExcludedError(RuntimeError): + pass + + +@dataclass +class CorpusGateSummary: + corpus: Optional[str] + kept_count: int = 0 + excluded_count: int = 0 + excluded_by_reason: Dict[str, List[str]] = field(default_factory=dict) + + @property + def total_count(self) -> int: + return self.kept_count + self.excluded_count + + @property + def excluded_ratio(self) -> float: + if not self.total_count: + return 0.0 + return self.excluded_count / self.total_count + + def __str__(self) -> str: + parts = [ + f'kept {self.kept_count} of {self.total_count} documents' + f' ({self.excluded_ratio:.0%} excluded)' + ] + for reason, document_ids in sorted(self.excluded_by_reason.items()): + parts.append(f'{reason}: {sorted(document_ids)}') + return '; '.join(parts) + + +def get_gate_summary_by_corpus( + verdict_and_corpus_list: Sequence[Any] +) -> Dict[Optional[str], CorpusGateSummary]: + """Summarise verdicts per corpus, from (verdict, corpus) pairs.""" + summary_by_corpus: Dict[Optional[str], CorpusGateSummary] = {} + for verdict, corpus in verdict_and_corpus_list: + summary = summary_by_corpus.setdefault(corpus, CorpusGateSummary(corpus=corpus)) + if not verdict.is_excluded: + summary.kept_count += 1 + continue + summary.excluded_count += 1 + assert verdict.primary_reason is not None + summary.excluded_by_reason.setdefault(verdict.primary_reason, []).append( + verdict.document_id + ) + return summary_by_corpus + + +def check_corpus_loss_or_fail( + summary_by_corpus: Mapping[Optional[str], CorpusGateSummary], + max_excluded_ratio: float +) -> None: + """Refuse rather than proceed when a corpus loses more than the stated share.""" + mostly_excluded = [ + summary + for summary in summary_by_corpus.values() + if summary.excluded_ratio > max_excluded_ratio + ] + if not mostly_excluded: + return + raise CorpusMostlyExcludedError( + 'excluded more than the configured %.0f%% of a corpus: %s' % ( + max_excluded_ratio * 100, + '; '.join( + f'{summary.corpus or "corpus not known"} {summary}' + for summary in mostly_excluded + ) + ) + ) diff --git a/tests/test_utils.py b/tests/test_utils.py index b220a2ac..ef06d0ff 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -31,7 +31,8 @@ def log_on_exception(f: T_WrappedCallableOrType) -> T_WrappedCallableOrType: @wraps(f) def wrapper(*args, **kwargs): try: - f(*args, **kwargs) + # returned, since wrapping a class wraps its helpers as well as its tests + return f(*args, **kwargs) except Exception as e: # pylint: disable=broad-except LOGGER.exception('failed due to %s', repr(e)) raise diff --git a/tests/training/cli/generate_delft_data_test.py b/tests/training/cli/generate_delft_data_test.py index 668157c5..4c756ff8 100644 --- a/tests/training/cli/generate_delft_data_test.py +++ b/tests/training/cli/generate_delft_data_test.py @@ -33,6 +33,7 @@ translate_tag_result_tags_IOB_to_grobid, translate_tags_IOB_to_grobid ) +from sciencebeam_parser.training.quality.gate import CorpusMostlyExcludedError from sciencebeam_parser.training.grobid_column_layout import ( get_grobid_column_layout_for_model_name ) @@ -733,3 +734,127 @@ def test_should_fall_back_when_the_file_does_not_carry_the_suffix(self): assert get_document_id_for_tei_file( '/tei/PPR459453.something-else.tei.xml', '.references.tei.xml' ) == 'PPR459453' + + +@log_on_exception +class TestQualityFilter: + def _write_tei(self, tei_source_path: Path, document_id: str, bibl_count: int) -> None: + tei_source_path.mkdir(parents=True, exist_ok=True) + ( + tei_source_path / f'{document_id}.references.referenceSegmenter.tei.xml' + ).write_bytes(etree.tostring(E('tei', E('text', E('listBibl', *[ + child + for index in range(bibl_count) + for child in (E('bibl', f'reference{index}', E('lb')), '\n') + ]))))) + + def _write_generated_record( + self, record_path: Path, rows: Sequence[dict] + ) -> None: + record_path.parent.mkdir(parents=True, exist_ok=True) + record_path.write_text( + '\n'.join(json.dumps(row) for row in rows) + '\n', encoding='utf-8' + ) + + def _run(self, tmp_path: Path, tei_source_path: Path, record_path: Path, *extra): + output_path = tmp_path / 'output.data' + main([ + '--model-name=reference_segmenter', + f'--tei-source-path={tei_source_path}/*.tei.xml', + f'--quality-record-path={record_path}', + f'--delft-output-path={output_path}', + '--quality-filter', + *extra + ]) + return output_path + + def test_should_leave_out_a_document_short_at_the_tei_stage(self, tmp_path: Path): + tei_source_path = tmp_path / 'train' / 'ore' / 'reference-segmenter' / 'corpus' / 'tei' + self._write_tei(tei_source_path, 'truncated', bibl_count=2) + for index in range(9): + self._write_tei(tei_source_path, f'sound{index}', bibl_count=4) + record_path = tmp_path / 'train' / 'ore' / 'reference-segmenter' / 'quality.jsonl' + self._write_generated_record(record_path, [ + { + 'document_id': 'truncated', 'written': True, + 'jats': {'status': 'ok', 'reference_count': 45}, + 'entity_element_count': 2, + }, + ] + [ + { + 'document_id': f'sound{index}', 'written': True, + 'jats': {'status': 'ok', 'reference_count': 4}, + 'entity_element_count': 4, + } + for index in range(9) + ]) + output_path = self._run(tmp_path, tei_source_path, record_path) + texts, _labels, _features = load_data_and_labels_crf_file(str(output_path)) + assert len(texts) == 9 + rows = { + json.loads(line)['document_id']: json.loads(line) + for line in Path( + str(output_path) + '.quality.jsonl' + ).read_text(encoding='utf-8').splitlines() + } + assert rows['truncated']['excluded'] is True + assert rows['truncated']['exclusion_reasons'] == ['elements-short-of-jats'] + assert rows['truncated']['exclusion_detail']['element_ratio'] == 0.044 + assert rows['sound0']['excluded'] is False + + def test_should_keep_every_document_without_the_filter(self, tmp_path: Path): + tei_source_path = tmp_path / 'train' / 'ore' / 'reference-segmenter' / 'corpus' / 'tei' + self._write_tei(tei_source_path, 'truncated', bibl_count=2) + record_path = tmp_path / 'train' / 'ore' / 'reference-segmenter' / 'quality.jsonl' + self._write_generated_record(record_path, [{ + 'document_id': 'truncated', 'written': True, + 'jats': {'status': 'ok', 'reference_count': 45}, + 'entity_element_count': 2, + }]) + output_path = tmp_path / 'output.data' + main([ + '--model-name=reference_segmenter', + f'--tei-source-path={tei_source_path}/*.tei.xml', + f'--quality-record-path={record_path}', + f'--delft-output-path={output_path}' + ]) + texts, _labels, _features = load_data_and_labels_crf_file(str(output_path)) + assert len(texts) == 1 + row = json.loads( + Path(str(output_path) + '.quality.jsonl').read_text(encoding='utf-8') + ) + assert 'excluded' not in row + + def test_should_refuse_when_a_corpus_is_mostly_excluded(self, tmp_path: Path): + tei_source_path = tmp_path / 'train' / 'ore' / 'reference-segmenter' / 'corpus' / 'tei' + for index in range(3): + self._write_tei(tei_source_path, f'truncated{index}', bibl_count=2) + record_path = tmp_path / 'train' / 'ore' / 'reference-segmenter' / 'quality.jsonl' + self._write_generated_record(record_path, [ + { + 'document_id': f'truncated{index}', 'written': True, + 'jats': {'status': 'ok', 'reference_count': 45}, + 'entity_element_count': 2, + } + for index in range(3) + ]) + with pytest.raises(CorpusMostlyExcludedError): + self._run(tmp_path, tei_source_path, record_path) + + def test_should_allow_a_stated_larger_loss(self, tmp_path: Path): + tei_source_path = tmp_path / 'train' / 'ore' / 'reference-segmenter' / 'corpus' / 'tei' + for index in range(3): + self._write_tei(tei_source_path, f'truncated{index}', bibl_count=2) + record_path = tmp_path / 'train' / 'ore' / 'reference-segmenter' / 'quality.jsonl' + self._write_generated_record(record_path, [ + { + 'document_id': f'truncated{index}', 'written': True, + 'jats': {'status': 'ok', 'reference_count': 45}, + 'entity_element_count': 2, + } + for index in range(3) + ]) + output_path = self._run( + tmp_path, tei_source_path, record_path, '--max-excluded-ratio=1.0' + ) + assert output_path.exists() diff --git a/tests/training/quality/gate_test.py b/tests/training/quality/gate_test.py new file mode 100644 index 00000000..0c4287c2 --- /dev/null +++ b/tests/training/quality/gate_test.py @@ -0,0 +1,223 @@ +import pytest + +from sciencebeam_parser.training.quality.gate import ( + CorpusMostlyExcludedError, + ExclusionReason, + ModelQualityThresholds, + QualityVerdict, + check_corpus_loss_or_fail, + get_gate_summary_by_corpus, + get_quality_verdict, + load_training_quality_config +) + + +REFERENCE_SEGMENTER = 'reference-segmenter' +DOCUMENT_ID_1 = 'document1' + +REFERENCE_SEGMENTER_THRESHOLDS = ModelQualityThresholds( + model_name=REFERENCE_SEGMENTER, + min_jats_reference_count=1, + min_element_ratio=0.9, + min_entity_ratio=0.8, +) + + +def _verdict(**kwargs) -> QualityVerdict: + return get_quality_verdict( + document_id=kwargs.pop('document_id', DOCUMENT_ID_1), + thresholds=kwargs.pop('thresholds', REFERENCE_SEGMENTER_THRESHOLDS), + **kwargs + ) + + +class TestLoadTrainingQualityConfig: + def test_should_load_the_shipped_thresholds(self): + config = load_training_quality_config() + thresholds = config.get_thresholds_for_model(REFERENCE_SEGMENTER) + assert thresholds.min_element_ratio == 0.9 + assert thresholds.min_entity_ratio == 0.8 + assert thresholds.has_cardinality_check + assert 0 < config.max_excluded_ratio < 1 + + def test_should_declare_a_region_model_as_having_no_cardinality(self): + thresholds = load_training_quality_config().get_thresholds_for_model('segmentation') + assert not thresholds.has_cardinality_check + assert thresholds.reason + + def test_should_declare_citation_without_a_label_floor(self): + thresholds = load_training_quality_config().get_thresholds_for_model('citation') + assert thresholds.min_element_ratio == 0.9 + assert thresholds.min_entity_ratio is None + assert not thresholds.label_floors + + def test_should_fail_for_a_model_with_no_entry_rather_than_pass_it(self): + with pytest.raises(KeyError): + load_training_quality_config().get_thresholds_for_model('not-a-model') + + def test_should_have_an_entry_for_every_model_generation_writes(self): + # A missing entry is a decision not taken; the gate must not pass it silently. + config = load_training_quality_config() + for model_name in [ + 'segmentation', 'header', 'affiliation-address', 'name-header', + 'name-citation', 'fulltext', 'figure', 'table', 'reference-segmenter', + 'citation', + ]: + assert config.get_thresholds_for_model(model_name) + + +class TestGetQualityVerdict: + def test_should_keep_a_document_that_is_right_throughout(self): + verdict = _verdict( + jats_status='ok', jats_reference_count=40, written=True, + entity_element_count=40, entity_start_count=40, sequence_count=1, + ) + assert not verdict.is_excluded + assert 'kept' in str(verdict) + + def test_should_exclude_a_document_whose_jats_declares_no_references(self): + verdict = _verdict( + jats_status='ok', jats_reference_count=0, written=False, sequence_count=1, + ) + assert verdict.primary_reason == ExclusionReason.JATS_HAS_NO_REFERENCES + assert verdict.detail['jats_reference_count'] == 0 + + def test_should_exclude_a_document_whose_jats_could_not_be_read(self): + verdict = _verdict( + jats_status='unparsable', written=True, + entity_element_count=0, entity_start_count=0, sequence_count=29, + ) + assert verdict.primary_reason == ExclusionReason.JATS_NOT_READABLE + assert verdict.detail['jats_status'] == 'unparsable' + + def test_should_exclude_a_document_short_at_the_tei_stage_naming_that_stage(self): + # PPR459453: a 45-entry reference list truncated to a 109-word region. + verdict = _verdict( + jats_status='ok', jats_reference_count=45, written=True, + entity_element_count=2, entity_start_count=2, sequence_count=1, + ) + assert verdict.exclusion_reasons == [ExclusionReason.ELEMENTS_SHORT_OF_JATS] + assert verdict.detail['element_ratio'] == 0.044 + + def test_should_exclude_a_document_collapsed_at_the_parse_naming_that_stage(self): + # The case that got through: correct in the TEI, one entity in the data. + verdict = _verdict( + jats_status='ok', jats_reference_count=40, written=True, + entity_element_count=40, entity_start_count=1, sequence_count=1, + ) + assert verdict.exclusion_reasons == [ExclusionReason.ENTITIES_SHORT_OF_ELEMENTS] + assert verdict.detail['entity_ratio'] == 0.025 + + def test_should_keep_a_document_holding_more_elements_than_the_jats_has(self): + # PPR534793: 22 references, 25 elements, from references split across a page. + verdict = _verdict( + jats_status='ok', jats_reference_count=22, written=True, + entity_element_count=25, entity_start_count=25, sequence_count=1, + ) + assert not verdict.is_excluded + + def test_should_keep_a_document_losing_one_entity_of_forty(self): + # An element holding nothing but a <label> cannot produce an entity. + verdict = _verdict( + jats_status='ok', jats_reference_count=37, written=True, + entity_element_count=37, entity_start_count=36, sequence_count=1, + ) + assert not verdict.is_excluded + + def test_should_exclude_a_document_with_no_training_sequences(self): + verdict = _verdict( + jats_status='ok', jats_reference_count=40, written=True, + entity_element_count=0, entity_start_count=0, sequence_count=0, + ) + assert ExclusionReason.NO_TRAINING_SEQUENCES in verdict.exclusion_reasons + + def test_should_name_the_earliest_stage_first(self): + verdict = _verdict( + jats_status='ok', jats_reference_count=0, written=False, sequence_count=0, + ) + assert verdict.primary_reason == ExclusionReason.JATS_HAS_NO_REFERENCES + + def test_should_not_treat_an_unavailable_count_as_a_failure(self): + # Assembly is run over corpora with no record at all. + verdict = _verdict(entity_start_count=40, sequence_count=1) + assert not verdict.is_excluded + + def test_should_apply_no_cardinality_check_to_a_region_model(self): + verdict = _verdict( + thresholds=ModelQualityThresholds( + model_name='segmentation', cardinality='none', reason='regions' + ), + jats_status='ok', jats_reference_count=0, written=True, sequence_count=1, + ) + assert not verdict.is_excluded + + +class TestGetGateSummaryByCorpus: + def test_should_count_what_it_kept_and_dropped_per_corpus(self): + summary_by_corpus = get_gate_summary_by_corpus([ + (QualityVerdict('kept1'), 'ore'), + (QualityVerdict('kept2'), 'ore'), + ( + QualityVerdict('dropped', [ExclusionReason.JATS_NOT_READABLE]), + 'ore' + ), + (QualityVerdict('kept3'), 'scielo_preprints-jats'), + ]) + assert summary_by_corpus['ore'].kept_count == 2 + assert summary_by_corpus['ore'].excluded_count == 1 + assert summary_by_corpus['scielo_preprints-jats'].excluded_count == 0 + + def test_should_report_the_dropped_documents_by_reason(self): + summary = get_gate_summary_by_corpus([ + ( + QualityVerdict('truncated', [ExclusionReason.ELEMENTS_SHORT_OF_JATS]), + 'ore' + ), + (QualityVerdict('kept'), 'ore'), + ])['ore'] + assert summary.excluded_by_reason == { + ExclusionReason.ELEMENTS_SHORT_OF_JATS: ['truncated'] + } + assert 'truncated' in str(summary) + assert 'kept 1 of 2' in str(summary) + + +class TestCheckCorpusLossOrFail: + def test_should_proceed_when_a_corpus_keeps_most_of_its_documents(self): + check_corpus_loss_or_fail( + get_gate_summary_by_corpus( + [(QualityVerdict(f'kept{index}'), 'ore') for index in range(9)] + + [(QualityVerdict('dropped', [ExclusionReason.JATS_NOT_READABLE]), 'ore')] + ), + max_excluded_ratio=0.2 + ) + + def test_should_refuse_when_a_corpus_is_mostly_excluded(self): + with pytest.raises(CorpusMostlyExcludedError) as exc_info: + check_corpus_loss_or_fail( + get_gate_summary_by_corpus([ + ( + QualityVerdict( + f'dropped{index}', + [ExclusionReason.ENTITIES_SHORT_OF_ELEMENTS] + ), + 'ore' + ) + for index in range(6) + ] + [(QualityVerdict(f'kept{index}'), 'ore') for index in range(4)]), + max_excluded_ratio=0.2 + ) + assert 'ore' in str(exc_info.value) + + def test_should_refuse_for_one_corpus_even_when_another_is_sound(self): + with pytest.raises(CorpusMostlyExcludedError): + check_corpus_loss_or_fail( + get_gate_summary_by_corpus( + [(QualityVerdict('dropped', [ExclusionReason.JATS_NOT_READABLE]), 'ore')] + + [ + (QualityVerdict(f'kept{index}'), 'scielo_preprints-jats') + for index in range(10) + ] + ), + max_excluded_ratio=0.2 + ) From 0669e9327495e7a00e7c7eee2a5b45b7fa9d941b Mon Sep 17 00:00:00 2001 From: Daniel Ecer <d.ecer@elifesciences.org> Date: Thu, 20 Aug 2026 14:02:46 +0100 Subject: [PATCH 4/4] Only refuse a JATS that failed, and leave nothing behind when refusing Three defects in the gate, from reviewing it rather than running it on the two corpora that happen to exercise none of them. A JATS status of `missing` means none was matched, which is what generating without --source-xml-path records for every document. The gate read any status other than ok as a JATS that could not be used, so such a corpus failed document by document and was then refused whole, for a reason that was not true. Only unparsable and unreadable are a defect now: a JATS that was never there leaves nothing to check against, and what cannot be checked is not a failure. The data is written as documents are assembled, so a refused corpus had already been written by the time the corpus-loss check ran, and whatever read that path next would find what looks like a corpus - the silent shrinkage the refusal exists to prevent. The output is now removed when assembly refuses, and the quality record kept so the refusal can be accounted for. A remote path cannot be removed by the writer in use, and says so rather than implying it was. no-training-sequences was collected before the stages that precede it, so a document failing both an earlier stage and that one reported the later one as its primary reason, against what the verdict promises. Stages are built in order. Also asserted: the models the config declares as having no cardinality are the models the counting module declares the same way. Two lists of one fact drift, and the gate would then count a cardinality it never checks. --- .../training/cli/generate_delft_data.py | 40 ++++++++++++--- sciencebeam_parser/training/quality/gate.py | 17 +++++-- .../training/cli/generate_delft_data_test.py | 20 ++++++++ tests/training/quality/gate_test.py | 50 +++++++++++++++++++ 4 files changed, 115 insertions(+), 12 deletions(-) diff --git a/sciencebeam_parser/training/cli/generate_delft_data.py b/sciencebeam_parser/training/cli/generate_delft_data.py index eb0d0060..7d18ba35 100644 --- a/sciencebeam_parser/training/cli/generate_delft_data.py +++ b/sciencebeam_parser/training/cli/generate_delft_data.py @@ -49,6 +49,7 @@ ) from sciencebeam_parser.training.quality.gate import ( TRAINING_QUALITY_CONFIG_FILE, + CorpusMostlyExcludedError, TrainingQualityConfig, check_corpus_loss_or_fail, get_gate_summary_by_corpus, @@ -414,11 +415,30 @@ def get_assembled_document_record( ) +def discard_refused_output(delft_output_path: str) -> None: + """Leave nothing usable behind when assembly refuses. + + The data is written as documents are assembled, so by the time a corpus is + known to be mostly excluded the file exists and would look like a corpus to + whatever reads it next. A remote path cannot be removed here, and says so. + """ + if not os.path.exists(delft_output_path): + LOGGER.warning( + 'assembly refused; the output at %r could not be removed and' + ' holds only the documents that passed', + delft_output_path + ) + return + os.remove(delft_output_path) + LOGGER.warning('assembly refused; removed the partial output at %r', delft_output_path) + + def log_gate_summary( model_name: str, assembled_records: Sequence[AssembledDocumentRecord], quality_config: TrainingQualityConfig, max_excluded_ratio: Optional[float] = None, + delft_output_path: Optional[str] = None, ) -> None: summary_by_corpus = get_gate_summary_by_corpus([ (record.verdict, record.corpus) @@ -435,12 +455,17 @@ def log_gate_summary( '%s / %s: %s', corpus or 'corpus not known', get_canonical_model_name(model_name), summary ) - check_corpus_loss_or_fail( - summary_by_corpus, - max_excluded_ratio - if max_excluded_ratio is not None - else quality_config.max_excluded_ratio - ) + try: + check_corpus_loss_or_fail( + summary_by_corpus, + max_excluded_ratio + if max_excluded_ratio is not None + else quality_config.max_excluded_ratio + ) + except CorpusMostlyExcludedError: + if delft_output_path: + discard_refused_output(delft_output_path) + raise def log_assembly_summary( @@ -566,7 +591,8 @@ def generate_delft_training_data( # pylint: disable=too-many-locals if quality_config is not None: log_gate_summary( model_name, assembled_records, quality_config, - max_excluded_ratio=max_excluded_ratio + max_excluded_ratio=max_excluded_ratio, + delft_output_path=delft_output_path ) diff --git a/sciencebeam_parser/training/quality/gate.py b/sciencebeam_parser/training/quality/gate.py index 523deda0..08d6bb16 100644 --- a/sciencebeam_parser/training/quality/gate.py +++ b/sciencebeam_parser/training/quality/gate.py @@ -26,6 +26,11 @@ NO_CARDINALITY = 'none' +# A JATS that was matched and could not be used is a defect; one that was never +# matched leaves nothing to check against, which is not the same thing -- data is +# legitimately generated with no JATS at all. +JATS_STATUSES_THAT_FAILED = frozenset({'unparsable', 'unreadable'}) + class ExclusionReason: """Why a document's training data is not used, in the order the stages run.""" @@ -132,20 +137,20 @@ def get_quality_verdict( # pylint: disable=too-many-branches """ reasons: List[str] = [] detail: Dict[str, Any] = {} - if sequence_count == 0: - reasons.append(ExclusionReason.NO_TRAINING_SEQUENCES) if not thresholds.has_cardinality_check: + if sequence_count == 0: + reasons.append(ExclusionReason.NO_TRAINING_SEQUENCES) return QualityVerdict(document_id, reasons, detail) - if jats_status is not None and jats_status != 'ok': - reasons.insert(0, ExclusionReason.JATS_NOT_READABLE) + if jats_status in JATS_STATUSES_THAT_FAILED: + reasons.append(ExclusionReason.JATS_NOT_READABLE) detail['jats_status'] = jats_status elif ( thresholds.min_jats_reference_count is not None and jats_reference_count is not None and jats_reference_count < thresholds.min_jats_reference_count ): - reasons.insert(0, ExclusionReason.JATS_HAS_NO_REFERENCES) + reasons.append(ExclusionReason.JATS_HAS_NO_REFERENCES) detail['jats_reference_count'] = jats_reference_count if written is False: reasons.append(ExclusionReason.NO_GENERATED_OUTPUT) @@ -172,6 +177,8 @@ def get_quality_verdict( # pylint: disable=too-many-branches detail['entity_start_count'] = entity_start_count detail['entity_element_count'] = entity_element_count + if sequence_count == 0: + reasons.append(ExclusionReason.NO_TRAINING_SEQUENCES) return QualityVerdict(document_id, reasons, detail) diff --git a/tests/training/cli/generate_delft_data_test.py b/tests/training/cli/generate_delft_data_test.py index 4c756ff8..c3e50790 100644 --- a/tests/training/cli/generate_delft_data_test.py +++ b/tests/training/cli/generate_delft_data_test.py @@ -825,6 +825,26 @@ def test_should_keep_every_document_without_the_filter(self, tmp_path: Path): ) assert 'excluded' not in row + def test_should_leave_no_usable_data_behind_when_it_refuses(self, tmp_path: Path): + tei_source_path = tmp_path / 'train' / 'ore' / 'reference-segmenter' / 'corpus' / 'tei' + for index in range(3): + self._write_tei(tei_source_path, f'truncated{index}', bibl_count=2) + record_path = tmp_path / 'train' / 'ore' / 'reference-segmenter' / 'quality.jsonl' + self._write_generated_record(record_path, [ + { + 'document_id': f'truncated{index}', 'written': True, + 'jats': {'status': 'ok', 'reference_count': 45}, + 'entity_element_count': 2, + } + for index in range(3) + ]) + with pytest.raises(CorpusMostlyExcludedError): + self._run(tmp_path, tei_source_path, record_path) + # a training run reading the path must not find a corpus there + assert not (tmp_path / 'output.data').exists() + # the record stays, so the refusal can be accounted for + assert (tmp_path / 'output.data.quality.jsonl').exists() + def test_should_refuse_when_a_corpus_is_mostly_excluded(self, tmp_path: Path): tei_source_path = tmp_path / 'train' / 'ore' / 'reference-segmenter' / 'corpus' / 'tei' for index in range(3): diff --git a/tests/training/quality/gate_test.py b/tests/training/quality/gate_test.py index 0c4287c2..e25a0dbc 100644 --- a/tests/training/quality/gate_test.py +++ b/tests/training/quality/gate_test.py @@ -1,5 +1,6 @@ import pytest +from sciencebeam_parser.training.quality.counting import MODELS_WITHOUT_ENTITY_COUNT from sciencebeam_parser.training.quality.gate import ( CorpusMostlyExcludedError, ExclusionReason, @@ -221,3 +222,52 @@ def test_should_refuse_for_one_corpus_even_when_another_is_sound(self): ), max_excluded_ratio=0.2 ) + + +class TestGetQualityVerdictWithoutJats: + def test_should_not_exclude_a_document_generated_without_any_jats(self): + # Training data is legitimately generated with no --source-xml-path, which + # records the JATS as missing. There is nothing to check against, and + # excluding on it would refuse the whole corpus. + verdict = get_quality_verdict( + document_id=DOCUMENT_ID_1, thresholds=REFERENCE_SEGMENTER_THRESHOLDS, + jats_status='missing', written=True, entity_start_count=40, sequence_count=1, + ) + assert not verdict.is_excluded + + def test_should_still_exclude_a_jats_that_could_not_be_read(self): + for jats_status in ['unparsable', 'unreadable']: + verdict = get_quality_verdict( + document_id=DOCUMENT_ID_1, thresholds=REFERENCE_SEGMENTER_THRESHOLDS, + jats_status=jats_status, written=True, sequence_count=1, + ) + assert verdict.primary_reason == ExclusionReason.JATS_NOT_READABLE + + def test_should_name_the_earliest_stage_first_when_several_failed(self): + # No training sequences is the last stage, so a shortfall before it leads. + verdict = get_quality_verdict( + document_id=DOCUMENT_ID_1, thresholds=REFERENCE_SEGMENTER_THRESHOLDS, + jats_status='ok', jats_reference_count=45, written=True, + entity_element_count=2, entity_start_count=0, sequence_count=0, + ) + assert verdict.primary_reason == ExclusionReason.ELEMENTS_SHORT_OF_JATS + assert verdict.exclusion_reasons[-1] == ExclusionReason.NO_TRAINING_SEQUENCES + + +class TestConfiguredModelsAgainstCounting: + def test_should_declare_the_same_models_as_having_no_entity_count(self): + # Two lists of the same fact drift apart; the gate would then report a + # cardinality it never counts, or count one it does not gate. + config = load_training_quality_config() + without_cardinality = { + model_name + for model_name, thresholds in config.thresholds_by_model.items() + if not thresholds.has_cardinality_check + } + assert without_cardinality == set(MODELS_WITHOUT_ENTITY_COUNT) + + def test_should_have_a_reason_for_every_model_without_a_check(self): + config = load_training_quality_config() + for thresholds in config.thresholds_by_model.values(): + if not thresholds.has_cardinality_check: + assert thresholds.reason, thresholds.model_name