diff --git a/doc/training.md b/doc/training.md index 671e3ddc..913a9d95 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. @@ -94,6 +146,87 @@ 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. + +#### 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_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/cli/generate_delft_data.py b/sciencebeam_parser/training/cli/generate_delft_data.py index e4dc102e..7d18ba35 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,29 @@ 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.gate import ( + TRAINING_QUALITY_CONFIG_FILE, + CorpusMostlyExcludedError, + 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, + 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 +94,50 @@ 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( + '--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', @@ -204,14 +271,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 +330,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 +338,158 @@ 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], + 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=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 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) + 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 + ) + 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( + 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 +498,12 @@ 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, + 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, @@ -309,23 +532,68 @@ 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 + ) + 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( delft_output_path, 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') - data_fp.writelines(iter_generate_delft_training_data_lines_for_document( + 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, training_tei_parser=training_tei_parser, data_generator=data_generator, column_layout=column_layout, include_extra_columns=include_extra_columns - )) + ) + 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 + ) + 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, + delft_output_path=delft_output_path + ) def run(args: argparse.Namespace): @@ -340,16 +608,24 @@ 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, + quality_filter=args.quality_filter, + quality_config_path=args.quality_config_path, + max_excluded_ratio=args.max_excluded_ratio ) 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/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/assembly.py b/sciencebeam_parser/training/quality/assembly.py new file mode 100644 index 00000000..672d5149 --- /dev/null +++ b/sciencebeam_parser/training/quality/assembly.py @@ -0,0 +1,214 @@ +"""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.training.quality.gate import QualityVerdict +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 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. + + 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 + verdict: Optional['QualityVerdict'] = 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) + 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 + + +@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 new file mode 100644 index 00000000..17c5a8b5 --- /dev/null +++ b/sciencebeam_parser/training/quality/counting.py @@ -0,0 +1,184 @@ +"""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, Iterator, Mapping, Optional, Sequence + +from lxml import etree + +from sciencebeam_parser.models.data import LabeledLayoutToken, 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', +}) + + +# 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): + 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.""" + 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 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 + ) + 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 + + +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/sciencebeam_parser/training/quality/gate.py b/sciencebeam_parser/training/quality/gate.py new file mode 100644 index 00000000..08d6bb16 --- /dev/null +++ b/sciencebeam_parser/training/quality/gate.py @@ -0,0 +1,254 @@ +"""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' + +# 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.""" + 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 not thresholds.has_cardinality_check: + if sequence_count == 0: + reasons.append(ExclusionReason.NO_TRAINING_SEQUENCES) + return QualityVerdict(document_id, reasons, detail) + + 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.append(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 + + if sequence_count == 0: + reasons.append(ExclusionReason.NO_TRAINING_SEQUENCES) + 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/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/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_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/cli/generate_delft_data_test.py b/tests/training/cli/generate_delft_data_test.py index 2ba7d947..c3e50790 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,10 +28,12 @@ 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 ) +from sciencebeam_parser.training.quality.gate import CorpusMostlyExcludedError from sciencebeam_parser.training.grobid_column_layout import ( get_grobid_column_layout_for_model_name ) @@ -612,3 +615,266 @@ 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' + + +@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_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): + 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/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/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 new file mode 100644 index 00000000..d8c9c199 --- /dev/null +++ b/tests/training/quality/counting_test.py @@ -0,0 +1,206 @@ +from lxml import etree + +from sciencebeam_parser.models.citation.labels import IDENTIFIER_LABEL +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, + count_entity_starts, + count_label_starts_per_sequence, + get_labels_for_model_data_list, + is_model_counted_by_label +) + + +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} + + +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') diff --git a/tests/training/quality/gate_test.py b/tests/training/quality/gate_test.py new file mode 100644 index 00000000..e25a0dbc --- /dev/null +++ b/tests/training/quality/gate_test.py @@ -0,0 +1,273 @@ +import pytest + +from sciencebeam_parser.training.quality.counting import MODELS_WITHOUT_ENTITY_COUNT +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 + ) + + +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 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