diff --git a/doc/training.md b/doc/training.md index b3923399..671e3ddc 100644 --- a/doc/training.md +++ b/doc/training.md @@ -80,6 +80,20 @@ It will do one of the following: - For models with only `tei` XML files (no layout feature), it will parse the `tei` and generate data using the data generator. - For models with additional layout data files, it will align the parsed `tei` with the layout data file and add the label to it. +The output matches GROBID's column layout for the model, so it can be mixed with +GROBID's own corpus. The expected layout per model is recorded in +[`grobid_column_layout.yml`](../sciencebeam_parser/resources/grobid_column_layout.yml), +and generating data for a model with no entry there fails rather than guessing. +`python -m sciencebeam_parser.training.cli.check_grobid_column_layout` re-checks +that file against GROBID's published corpora; it downloads them, so it is run by +hand rather than in CI. + +Pass `--include-extra-columns` to also emit the columns this project adds on top +of GROBID's layout. Today that is the `segmentation` model's `whole_line_text`, +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. + #### Example command for `segmentation` model ```bash diff --git a/pyproject.toml b/pyproject.toml index a26ac83f..aa08b970 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -65,6 +65,7 @@ include = ["sciencebeam_parser*"] [tool.setuptools.package-data] "sciencebeam_parser" = [ "py.typed", + "resources/*.yml", "resources/default_config/*.yml", "resources/xslt/*.xsl" ] diff --git a/sciencebeam_parser/resources/grobid_column_layout.yml b/sciencebeam_parser/resources/grobid_column_layout.yml new file mode 100644 index 00000000..8142e793 --- /dev/null +++ b/sciencebeam_parser/resources/grobid_column_layout.yml @@ -0,0 +1,287 @@ +# The column layout of GROBID's own training data, per model. +# +# GROBID's FeaturesVector*.printVector() ends with a label slot: +# +# if (label != null) +# res.append(" " + label + "\n"); +# else +# res.append(" 0\n"); +# +# label_slot records what GROBID's corpus writer does with it, which decides +# what a training data line looks like: +# +# filled the writer sets features.label = tag, so the label occupies the +# slot: token, columns, label +# unfilled the writer leaves it null and appends the tag, so the slot stays +# `0`: token, columns, 0, label +# absent printVector has the block commented out: token, columns, label +# +# columns are the feature columns excluding the slot, starting with the token +# itself, and match the data generator's feature_names position by position. +# extra_columns are columns we add that GROBID has no counterpart for; they are +# emitted only with --include-extra-columns. +# +# reference_training_corpus is the labelled corpus the recorded layout was +# measured against, and what +# `python -m sciencebeam_parser.training.cli.check_grobid_column_layout` reads. +# It is training-specific: the inference vector is columns plus extra_columns +# plus the slot, and has no published counterpart here. +# +# Models sharing a data generator share an entry, via a YAML anchor. + +models: + + # whole_line_text is ours: GROBID has no counterpart for it, and it sits at + # the index GROBID uses for the label. delft models read it as a text + # feature; wapiti templates stop before it. + segmentation: + generator: SegmentationDataGenerator + label_slot: absent + reference_training_corpus: + - https://github.com/eLifePathways/sciencebeam-datasets/releases/download/grobid-0.9.0/delft-grobid-0.9.0-segmentation.train.gz + columns: + - token_text + - second_token_text + - lower_token_text + - prefix_1 + - prefix_2 + - prefix_3 + - prefix_4 + - block_status + - page_status + - token_font_status + - token_font_size + - is_bold + - is_italic + - capitalisation + - digit_status + - is_single_char + - is_proper_name + - is_common_name + - is_first_name + - is_year + - is_month + - is_email + - is_http + - relative_document_position + - relative_page_position + - punctuation_profile + - punctuation_profile_length + - block_relative_line_length + - is_bitmap_around + - is_vector_around + - is_repetitive_pattern + - is_first_repetitive_pattern + - is_main_area + extra_columns: + - whole_line_text + + header: + generator: HeaderDataGenerator + label_slot: unfilled + reference_training_corpus: + - https://github.com/eLifePathways/sciencebeam-datasets/releases/download/grobid-0.9.0/delft-grobid-0.9.0-header.train.gz + columns: + - token_text + - lower_token_text + - prefix_1 + - prefix_2 + - prefix_3 + - prefix_4 + - suffix_1 + - suffix_2 + - suffix_3 + - suffix_4 + - block_status + - line_status + - alignment + - token_font_status + - token_font_size + - is_bold + - is_italic + - capitalisation + - digit_status + - is_single_char + - is_proper_name + - is_common_name + - is_year + - is_month + - is_location_name + - is_email + - is_http + - punctuation_type + - is_largest_font + - is_smallest_font + - is_larger_than_average_font + + # figure and table share this layout: GROBID generates their training data + # with the fulltext feature vector (FeaturesVectorFigure is unreferenced), so + # the fulltext corpus is the reference for all three. Their own published + # corpora are not: at 0.9.0 the figure one is 26 feature columns throughout + # and the table one mixes 26 across 11569 lines with 27 across 5837, the + # 27-column lines being those that carry is_superscript. + fulltext: &fulltext + generator: FullTextDataGenerator + label_slot: absent + reference_training_corpus: + - https://github.com/eLifePathways/sciencebeam-datasets/releases/download/grobid-0.9.0/delft-grobid-0.9.0-fulltext.train.gz + columns: + - token_text + - lower_token_text + - prefix_1 + - prefix_2 + - prefix_3 + - prefix_4 + - suffix_1 + - suffix_2 + - suffix_3 + - suffix_4 + - block_status + - line_status + - alignment + - token_font_status + - token_font_size + - is_bold + - is_italic + - capitalisation + - digit_status + - is_single_char + - punctuation_type + - relative_document_position + - relative_page_position + - is_bitmap_around + - callout_type + - is_callout_known + - is_superscript + + figure: *fulltext + + table: *fulltext + + reference_segmenter: + generator: ReferenceSegmenterDataGenerator + label_slot: unfilled + reference_training_corpus: + - https://github.com/eLifePathways/sciencebeam-datasets/releases/download/grobid-0.9.0/delft-grobid-0.9.0-reference-segmenter.train.gz + columns: + - token_text + - lower_token_text + - prefix_1 + - prefix_2 + - prefix_3 + - prefix_4 + - suffix_1 + - suffix_2 + - suffix_3 + - suffix_4 + - line_status + - alignment + - capitalisation + - digit_status + - is_single_char + - is_proper_name + - is_common_name + - is_first_name + - is_location_name + - is_year + - is_month + - is_http + - punctuation_profile + - line_token_relative_position + - line_relative_length + - block_status + - truncated_punctuation_profile_length + + citation: + generator: CitationDataGenerator + label_slot: filled + reference_training_corpus: + - https://github.com/eLifePathways/sciencebeam-datasets/releases/download/grobid-0.9.0/delft-grobid-0.9.0-citation.train.gz + columns: + - token_text + - lower_token_text + - prefix_1 + - prefix_2 + - prefix_3 + - prefix_4 + - suffix_1 + - suffix_2 + - suffix_3 + - suffix_4 + - line_status + - capitalisation + - digit_status + - is_single_char + - is_proper_name + - is_common_name + - is_first_name + - is_last_name + - is_location_name + - is_year + - is_month + - is_http + - is_known_collaboration + - is_known_journal_title + - is_known_conference_title + - is_known_publisher + - is_known_identifier + - punctuation_type + - sentence_token_relative_position + + affiliation_address: + generator: AffiliationAddressDataGenerator + label_slot: filled + reference_training_corpus: + - https://github.com/eLifePathways/sciencebeam-datasets/releases/download/grobid-0.9.0/delft-grobid-0.9.0-affiliation-address.train.gz + columns: + - token_text + - lower_token_text + - prefix_1 + - prefix_2 + - prefix_3 + - prefix_4 + - suffix_1 + - suffix_2 + - suffix_3 + - suffix_4 + - line_status + - capitalisation + - digit_status + - is_single_char + - is_proper_name + - is_common_name + - is_first_name + - is_location_name + - is_country + - punctuation_type + - word_shape + + name_header: &name + generator: NameDataGenerator + label_slot: filled + reference_training_corpus: + - https://github.com/eLifePathways/sciencebeam-datasets/releases/download/grobid-0.9.0/delft-grobid-0.9.0-name-header.train.gz + - https://github.com/eLifePathways/sciencebeam-datasets/releases/download/grobid-0.9.0/delft-grobid-0.9.0-name-citation.train.gz + columns: + - token_text + - lower_token_text + - prefix_1 + - prefix_2 + - prefix_3 + - prefix_4 + - suffix_1 + - suffix_2 + - suffix_3 + - suffix_4 + - line_status + - capitalisation + - digit_status + - is_single_char + - is_common_name + - is_first_name + - is_last_name + - is_known_title + - is_known_suffix + - punctuation_type + + name_citation: *name diff --git a/sciencebeam_parser/training/cli/check_grobid_column_layout.py b/sciencebeam_parser/training/cli/check_grobid_column_layout.py new file mode 100644 index 00000000..a00b8931 --- /dev/null +++ b/sciencebeam_parser/training/cli/check_grobid_column_layout.py @@ -0,0 +1,149 @@ +"""Cross-check the recorded GROBID column layout against GROBID's own corpora. + +Run by hand: it downloads the reference corpora, so it is deliberately not part +of the test suite. The offline per-model test asserts that the data generators +still match what is recorded here; this asserts that what is recorded still +matches GROBID. +""" + +import argparse +import logging +from collections import Counter +from itertools import islice +from typing import Dict, List, Optional, Tuple + +from sciencebeam_trainer_delft.utils.io import auto_download_input_file + +from sciencebeam_parser.training.grobid_column_layout import ( + GrobidColumnLayout, + LabelSlot, + load_grobid_column_layout_by_name +) + + +LOGGER = logging.getLogger(__name__) + + +DEFAULT_MAX_LINES = 200000 + + +def parse_args(argv: Optional[List[str]] = None) -> argparse.Namespace: + parser = argparse.ArgumentParser( + 'ScienceBeam Parser: Check the recorded GROBID column layout' + ) + parser.add_argument( + '--model-name', + type=str, + action='append', + help='Model to check (repeatable, defaults to every recorded model)' + ) + parser.add_argument( + '--corpus', + type=str, + action='append', + help=( + 'Corpus to check against instead of the recorded' + ' reference_training_corpus (repeatable)' + ) + ) + parser.add_argument( + '--max-lines', + type=int, + default=DEFAULT_MAX_LINES, + help='Token lines to read per corpus' + ) + parser.add_argument( + '--debug', + action='store_true', + help='Enable debug logging' + ) + return parser.parse_args(argv) + + +def get_corpus_stats(filename: str, max_lines: int) -> Tuple[Counter, Counter]: + column_counts: Counter = Counter() + trailing_values: Counter = Counter() + with auto_download_input_file(filename, auto_decompress=True) as local_file: + with open(local_file, 'r', encoding='utf-8') as fp: + for line in islice((line for line in fp if line.strip()), max_lines): + columns = line.split() + column_counts[len(columns)] += 1 + if len(columns) >= 2: + trailing_values[columns[-2]] += 1 + return column_counts, trailing_values + + +def check_corpus( + layout: GrobidColumnLayout, + filename: str, + max_lines: int +) -> List[str]: + expected_column_count = len(layout.get_training_data_column_names()) + 1 + column_counts, trailing_values = get_corpus_stats(filename, max_lines) + problems: List[str] = [] + if not column_counts: + problems.append('no token lines found') + if set(column_counts) - {expected_column_count}: + problems.append( + 'expected %d columns, found %r' % ( + expected_column_count, dict(sorted(column_counts.items())) + ) + ) + if layout.label_slot == LabelSlot.UNFILLED and set(trailing_values) - {'0'}: + problems.append( + 'label_slot is %r, but the column before the label is not always `0`: %r' % ( + layout.label_slot, dict(trailing_values.most_common(5)) + ) + ) + LOGGER.info( + '%s: %s\n columns=%r\n column before the label=%r', + layout.name, filename, + dict(sorted(column_counts.items())), + dict(trailing_values.most_common(5)) + ) + return problems + + +def run(args: argparse.Namespace) -> int: + layout_by_name = load_grobid_column_layout_by_name() + model_names = args.model_name or sorted(layout_by_name) + if args.corpus and len(model_names) != 1: + raise ValueError('--corpus applies to a single --model-name') + problem_count = 0 + checked: Dict[Tuple[str, str, int], str] = {} + for model_name in model_names: + layout = layout_by_name[model_name] + corpus_list = args.corpus or list(layout.reference_training_corpus) + if not corpus_list: + LOGGER.warning('%s: no reference corpus recorded', model_name) + continue + for filename in corpus_list: + # models sharing a layout share a corpus; downloading it once is enough + key = (filename, layout.label_slot, len(layout.get_training_data_column_names())) + already_checked_for = checked.get(key) + if already_checked_for: + LOGGER.info('%s: same check as %s', model_name, already_checked_for) + continue + checked[key] = model_name + for problem in check_corpus(layout, filename, max_lines=args.max_lines): + problem_count += 1 + LOGGER.error('%s: %s: %s', model_name, filename, problem) + if problem_count: + LOGGER.error('%d problem(s) found', problem_count) + else: + LOGGER.info('recorded layout agrees with every corpus checked') + return 1 if problem_count else 0 + + +def main(argv: Optional[List[str]] = None) -> int: + args = parse_args(argv) + if args.debug: + for name in [__name__, 'sciencebeam_parser', 'sciencebeam_trainer_delft']: + logging.getLogger(name).setLevel('DEBUG') + return run(args) + + +if __name__ == '__main__': + logging.basicConfig(level='INFO') + + raise SystemExit(main()) diff --git a/sciencebeam_parser/training/cli/generate_delft_data.py b/sciencebeam_parser/training/cli/generate_delft_data.py index e30caec7..e4dc102e 100644 --- a/sciencebeam_parser/training/cli/generate_delft_data.py +++ b/sciencebeam_parser/training/cli/generate_delft_data.py @@ -32,6 +32,12 @@ ModelDataGenerator ) from sciencebeam_parser.models.training_data import TrainingTeiParser +from sciencebeam_parser.training.grobid_column_layout import ( + GrobidColumnLayout, + get_grobid_column_layout_for_model_name, + get_validated_training_data_feature_indices, + select_feature_columns +) from sciencebeam_parser.resources.default_config import DEFAULT_CONFIG_FILE from sciencebeam_parser.config.config import AppConfig @@ -65,6 +71,16 @@ def parse_args(argv: Optional[List[str]] = None) -> argparse.Namespace: type=str, required=True ) + parser.add_argument( + '--include-extra-columns', + action='store_true', + help=( + 'Emit the columns this project adds on top of GROBID\'s layout' + ' (segmentation\'s whole_line_text, read by delft models as a text feature).' + ' Without it the output matches GROBID\'s column layout and can be mixed' + ' with GROBID\'s own corpus.' + ) + ) parser.add_argument( '--debug', action='store_true', @@ -192,7 +208,9 @@ def iter_generate_delft_training_data_lines_for_document( # pylint: disable=too tei_file: str, raw_file: Optional[str], training_tei_parser: TrainingTeiParser, - data_generator: ModelDataGenerator + data_generator: ModelDataGenerator, + column_layout: GrobidColumnLayout, + include_extra_columns: bool = False ) -> Iterable[str]: with auto_download_input_file( tei_file, @@ -239,20 +257,30 @@ 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 + feature_indices = get_validated_training_data_feature_indices( + column_layout, + feature_column_count=len(features[0][0]), + data_generator_name=type(data_generator).__name__, + 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=features + features=select_feature_columns(features, feature_indices) ) -def generate_delft_training_data( +def generate_delft_training_data( # pylint: disable=too-many-locals model_name: str, tei_source_path: str, raw_source_path: str, delft_output_path: str, - sciencebeam_parser: ScienceBeamParser + sciencebeam_parser: ScienceBeamParser, + include_extra_columns: bool = False ): training_tei_parser = get_training_tei_parser_for_model_name( model_name, @@ -262,6 +290,12 @@ def generate_delft_training_data( model_name, sciencebeam_parser=sciencebeam_parser ) + column_layout = get_grobid_column_layout_for_model_name(model_name) + LOGGER.info( + 'column layout for %r: %d columns, label_slot=%r, extra_columns=%r (included: %r)', + model_name, len(column_layout.columns), column_layout.label_slot, + list(column_layout.extra_columns), include_extra_columns + ) LOGGER.debug('tei_source_path: %r', tei_source_path) tei_file_list = glob(tei_source_path) if not tei_file_list: @@ -288,7 +322,9 @@ def generate_delft_training_data( tei_file=tei_file, raw_file=raw_file, training_tei_parser=training_tei_parser, - data_generator=data_generator + data_generator=data_generator, + column_layout=column_layout, + include_extra_columns=include_extra_columns )) @@ -303,7 +339,8 @@ def run(args: argparse.Namespace): tei_source_path=args.tei_source_path, raw_source_path=args.raw_source_path, delft_output_path=args.delft_output_path, - sciencebeam_parser=sciencebeam_parser + sciencebeam_parser=sciencebeam_parser, + include_extra_columns=args.include_extra_columns ) diff --git a/sciencebeam_parser/training/grobid_column_layout.py b/sciencebeam_parser/training/grobid_column_layout.py new file mode 100644 index 00000000..78946073 --- /dev/null +++ b/sciencebeam_parser/training/grobid_column_layout.py @@ -0,0 +1,162 @@ +import os +from dataclasses import dataclass, field +from typing import Dict, List, Mapping, Optional, Sequence + +import numpy as np +import yaml + + +GROBID_COLUMN_LAYOUT_FILE = os.path.join( + os.path.dirname(os.path.dirname(os.path.abspath(__file__))), + 'resources', + 'grobid_column_layout.yml' +) + + +PLACEHOLDER_COLUMN_NAME = 'dummy_label' + + +class LabelSlot: + FILLED = 'filled' + UNFILLED = 'unfilled' + ABSENT = 'absent' + + +VALID_LABEL_SLOTS = frozenset({LabelSlot.FILLED, LabelSlot.UNFILLED, LabelSlot.ABSENT}) + +LABEL_SLOTS_WITH_PLACEHOLDER_COLUMN = frozenset({LabelSlot.FILLED, LabelSlot.UNFILLED}) + + +@dataclass(frozen=True) +class GrobidColumnLayout: + name: str + generator: str + label_slot: str + columns: Sequence[str] + extra_columns: Sequence[str] = field(default_factory=tuple) + reference_training_corpus: Sequence[str] = field(default_factory=tuple) + + def __post_init__(self): + if self.label_slot not in VALID_LABEL_SLOTS: + raise ValueError( + 'invalid label_slot for %r: %r (expected one of %s)' % ( + self.name, self.label_slot, sorted(VALID_LABEL_SLOTS) + ) + ) + + @property + def has_placeholder_column(self) -> bool: + return self.label_slot in LABEL_SLOTS_WITH_PLACEHOLDER_COLUMN + + def get_data_generator_column_names(self) -> List[str]: + """The columns a data generator emits, which is what inference gets.""" + names = list(self.columns) + list(self.extra_columns) + if self.has_placeholder_column: + names.append(PLACEHOLDER_COLUMN_NAME) + return names + + def get_training_data_column_names( + self, + include_extra_columns: bool = False + ) -> List[str]: + """The columns of a training data line, before the label.""" + names = list(self.columns) + if include_extra_columns: + names.extend(self.extra_columns) + if self.label_slot == LabelSlot.UNFILLED: + names.append(PLACEHOLDER_COLUMN_NAME) + return names + + def get_training_data_feature_indices( + self, + include_extra_columns: bool = False + ) -> List[int]: + """Indices into a feature row, which excludes the leading token column.""" + generator_column_names = self.get_data_generator_column_names()[1:] + return [ + generator_column_names.index(name) + for name in self.get_training_data_column_names(include_extra_columns)[1:] + ] + + +def _get_grobid_column_layout(name: str, layout_config: Mapping) -> GrobidColumnLayout: + return GrobidColumnLayout( + name=name, + generator=layout_config['generator'], + label_slot=layout_config['label_slot'], + columns=tuple(layout_config['columns']), + extra_columns=tuple(layout_config.get('extra_columns') or ()), + reference_training_corpus=tuple(layout_config.get('reference_training_corpus') or ()) + ) + + +def load_grobid_column_layout_by_name( + filename: str = GROBID_COLUMN_LAYOUT_FILE +) -> Dict[str, GrobidColumnLayout]: + with open(filename, 'r', encoding='utf-8') as fp: + config = yaml.safe_load(fp) + return { + name: _get_grobid_column_layout(name, layout_config) + for name, layout_config in config['models'].items() + } + + +def get_grobid_column_layout_for_model_name( + model_name: str, + filename: str = GROBID_COLUMN_LAYOUT_FILE +) -> GrobidColumnLayout: + layout_by_name = load_grobid_column_layout_by_name(filename) + layout = layout_by_name.get(model_name) + if layout is None: + raise ValueError( + 'no GROBID column layout recorded for model %r, add one to %s (known: %s)' % ( + model_name, os.path.basename(filename), sorted(layout_by_name) + ) + ) + return layout + + +def get_validated_training_data_feature_indices( + layout: GrobidColumnLayout, + feature_column_count: int, + data_generator_name: Optional[str] = None, + data_generator_column_names: Optional[Sequence[str]] = None, + include_extra_columns: bool = False +) -> List[int]: + if data_generator_name is not None and data_generator_name != layout.generator: + raise ValueError( + 'model %r uses data generator %r, but the recorded layout is for %r' % ( + layout.name, data_generator_name, layout.generator + ) + ) + expected_column_names = layout.get_data_generator_column_names() + if ( + data_generator_column_names is not None + and list(data_generator_column_names) != expected_column_names + ): + raise ValueError( + 'columns of %r do not match the recorded layout: %r != %r' % ( + layout.name, list(data_generator_column_names), expected_column_names + ) + ) + expected_feature_column_count = len(expected_column_names) - 1 + if feature_column_count != expected_feature_column_count: + raise ValueError( + 'expected %d feature columns for %r, but found %d' % ( + expected_feature_column_count, layout.name, feature_column_count + ) + ) + return layout.get_training_data_feature_indices(include_extra_columns) + + +def select_feature_columns( + features: np.ndarray, + feature_indices: Sequence[int] +) -> np.ndarray: + return np.asarray([ + [ + [token_features[index] for index in feature_indices] + for token_features in document_features + ] + for document_features in features.tolist() + ], dtype=object) diff --git a/tests/training/cli/generate_delft_data_test.py b/tests/training/cli/generate_delft_data_test.py index e98aff2b..2ba7d947 100644 --- a/tests/training/cli/generate_delft_data_test.py +++ b/tests/training/cli/generate_delft_data_test.py @@ -2,7 +2,7 @@ import logging import gzip from pathlib import Path -from typing import Iterator, Optional, Sequence +from typing import Iterator, List, Optional, Sequence from unittest.mock import MagicMock, patch import pytest @@ -31,6 +31,9 @@ translate_tag_result_tags_IOB_to_grobid, translate_tags_IOB_to_grobid ) +from sciencebeam_parser.training.grobid_column_layout import ( + get_grobid_column_layout_for_model_name +) from sciencebeam_parser.utils.xml_writer import XmlTreeWriter from tests.processors.fulltext.model_mocks import MockFullTextModels @@ -70,13 +73,40 @@ def _document_features_context( ) -def _test_generate_delft_with_multiple_tokens_tei_and_raw( +def _get_raw_feature_rows(model_name: str, token_count: int) -> List[List[str]]: + """Distinguishable placeholder values, one row per token, at the width the + data generator would have produced.""" + layout = get_grobid_column_layout_for_model_name(model_name) + column_count = len(layout.get_data_generator_column_names()) - 1 + return [ + [f'{token_index}.{column_index}' for column_index in range(column_count)] + for token_index in range(token_count) + ] + + +def _get_expected_training_feature_rows( + model_name: str, + raw_feature_rows: Sequence[Sequence[str]], + include_extra_columns: bool = False +) -> List[List[str]]: + layout = get_grobid_column_layout_for_model_name(model_name) + column_count = len( + layout.get_training_data_column_names(include_extra_columns) + ) - 1 + assert layout.get_training_data_feature_indices(include_extra_columns) == list( + range(column_count) + ), 'expected the emitted columns to be a prefix of the generated ones' + return [list(row[:column_count]) for row in raw_feature_rows] + + +def _test_generate_delft_with_multiple_tokens_tei_and_raw( # pylint: disable=too-many-locals tmp_path: Path, model_name: str, file_suffix: str, tei_root: etree.ElementBase, tokens: Sequence[str], - expected_labels: Sequence[str] + expected_labels: Sequence[str], + include_extra_columns: bool = False ): assert len(tokens) == len(expected_labels) tei_source_path = tmp_path / 'tei' @@ -87,21 +117,26 @@ def _test_generate_delft_with_multiple_tokens_tei_and_raw( etree.tostring(tei_root) ) raw_source_path.mkdir(parents=True, exist_ok=True) - expected_features = [[ - [f'{i}.1', f'{i}.2', f'{i}.3'] - for i in range(len(tokens)) - ]] + raw_feature_rows = _get_raw_feature_rows(model_name, len(tokens)) (raw_source_path / f'sample{file_suffix}').write_text('\n'.join([ - f'{token} {" ".join(expected_token_features)}' - for token, expected_token_features in zip(tokens, expected_features[0]) + f'{token} {" ".join(raw_token_features)}' + for token, raw_token_features in zip(tokens, raw_feature_rows) ])) main([ f'--model-name={model_name}', f'--tei-source-path={tei_source_path}/*.tei.xml', f'--raw-source-path={raw_source_path}', f'--delft-output-path={output_path}' - ]) + ] + (['--include-extra-columns'] if include_extra_columns else [])) assert output_path.exists() + expected_features = [_get_expected_training_feature_rows( + model_name, raw_feature_rows, include_extra_columns + )] + assert [ + len(line.split()) + for line in output_path.read_text().splitlines() + if line.strip() + ] == [1 + len(expected_features[0][0]) + 1] * len(tokens) texts, _labels, _features = load_data_and_labels_crf_file( str(output_path) ) @@ -117,7 +152,8 @@ def _test_generate_delft_with_two_tokens_tei_and_raw( model_name: str, file_suffix: str, tei_root: etree.ElementBase, - expected_labels: Sequence[str] + expected_labels: Sequence[str], + include_extra_columns: bool = False ): _test_generate_delft_with_multiple_tokens_tei_and_raw( tmp_path=tmp_path, @@ -125,7 +161,8 @@ def _test_generate_delft_with_two_tokens_tei_and_raw( file_suffix=file_suffix, tei_root=tei_root, expected_labels=expected_labels, - tokens=[TOKEN_1, TOKEN_2] + tokens=[TOKEN_1, TOKEN_2], + include_extra_columns=include_extra_columns ) @@ -158,7 +195,10 @@ def _test_generate_delft_with_multiple_tokens_tei_only( # pylint: disable=too-m expected_data_lines = list(data_generator.iter_data_lines_for_layout_document( layout_document )) - _expected_texts, expected_features = load_data_crf_lines(expected_data_lines) + _expected_texts, generated_features = load_data_crf_lines(expected_data_lines) + expected_features = [_get_expected_training_feature_rows( + model_name, generated_features.tolist()[0] + )] LOGGER.debug('expected_features: %r', expected_features) texts, labels, features = load_data_and_labels_crf_file( str(output_path) @@ -170,7 +210,7 @@ def _test_generate_delft_with_multiple_tokens_tei_only( # pylint: disable=too-m assert len(texts) == 1 assert list(texts[0]) == tokens assert list(labels[0]) == expected_labels - assert features.tolist() == expected_features.tolist() + assert features.tolist() == expected_features def _test_generate_delft_with_two_tokens_tei_only( @@ -249,6 +289,52 @@ def test_should_be_able_to_generate_segmentation_training_data( expected_labels=['B-
', 'B-'] ) + def test_should_include_the_segmentation_extra_column_when_asked_for( + self, + tmp_path: Path + ): + _test_generate_delft_with_two_tokens_tei_and_raw( + tmp_path=tmp_path, + model_name='segmentation', + file_suffix='.segmentation', + tei_root=E('tei', E('text', *[ + E('front', TOKEN_1, E('lb')), + '\n', + E('body', TOKEN_2, E('lb')), + '\n' + ])), + expected_labels=['B-
', 'B-'], + include_extra_columns=True + ) + + def test_should_reject_a_raw_file_of_another_width( + self, + tmp_path: Path + ): + tei_source_path = tmp_path / 'tei' + raw_source_path = tmp_path / 'raw' + tei_source_path.mkdir(parents=True, exist_ok=True) + (tei_source_path / 'sample.segmentation.tei.xml').write_bytes(etree.tostring( + E('tei', E('text', *[ + E('front', TOKEN_1, E('lb')), + '\n', + E('body', TOKEN_2, E('lb')), + '\n' + ])) + )) + raw_source_path.mkdir(parents=True, exist_ok=True) + (raw_source_path / 'sample.segmentation').write_text('\n'.join([ + f'{TOKEN_1} feature1 feature2', + f'{TOKEN_2} feature1 feature2' + ])) + with pytest.raises(ValueError): + main([ + '--model-name=segmentation', + f'--tei-source-path={tei_source_path}/*.tei.xml', + f'--raw-source-path={raw_source_path}', + f'--delft-output-path={tmp_path}/output.data' + ]) + def test_should_be_able_to_generate_header_training_data( self, tmp_path: Path @@ -508,13 +594,10 @@ def test_should_be_able_to_load_and_generate_gzipped_training_data( gzip.compress(etree.tostring(tei_root)) ) raw_source_path.mkdir(parents=True, exist_ok=True) - expected_features = [[ - [f'{i}.1', f'{i}.2', f'{i}.3'] - for i in range(len(tokens)) - ]] + raw_feature_rows = _get_raw_feature_rows(model_name, len(tokens)) (raw_source_path / f'sample{file_suffix}.gz').write_text('\n'.join([ - f'{token} {" ".join(expected_token_features)}' - for token, expected_token_features in zip(tokens, expected_features[0]) + f'{token} {" ".join(raw_token_features)}' + for token, raw_token_features in zip(tokens, raw_feature_rows) ])) main([ f'--model-name={model_name}', diff --git a/tests/training/grobid_column_layout_test.py b/tests/training/grobid_column_layout_test.py new file mode 100644 index 00000000..3f9c50f2 --- /dev/null +++ b/tests/training/grobid_column_layout_test.py @@ -0,0 +1,267 @@ +import dataclasses +import logging +from typing import Dict, Iterable + +import pytest + +from sciencebeam_parser.models.data import ( + DEFAULT_DOCUMENT_FEATURES_CONTEXT, + ModelDataGenerator +) +from sciencebeam_parser.models.model import Model +from sciencebeam_parser.processors.fulltext.models import FullTextModels +from sciencebeam_parser.training.grobid_column_layout import ( + GrobidColumnLayout, + LabelSlot, + PLACEHOLDER_COLUMN_NAME, + get_grobid_column_layout_for_model_name, + get_validated_training_data_feature_indices, + load_grobid_column_layout_by_name +) + +from tests.processors.fulltext.model_mocks import MockFullTextModels + + +LOGGER = logging.getLogger(__name__) + + +# The column count of a training data line, i.e. the token, the feature columns +# and the label. Recorded independently of the layout file, so that editing the +# layout without also changing GROBID's corpus fails here. +EXPECTED_TRAINING_DATA_COLUMN_COUNT_BY_MODEL_NAME = { + 'segmentation': 34, + 'header': 33, + 'fulltext': 28, + 'figure': 28, + 'table': 28, + 'reference_segmenter': 29, + 'citation': 30, + 'affiliation_address': 22, + 'name_header': 21, + 'name_citation': 21 +} + + +EXPECTED_TRAINING_DATA_COLUMN_COUNT_WITH_EXTRA_COLUMNS_BY_MODEL_NAME = { + **EXPECTED_TRAINING_DATA_COLUMN_COUNT_BY_MODEL_NAME, + 'segmentation': 35 +} + + +LAYOUT_BY_MODEL_NAME = load_grobid_column_layout_by_name() + +MODEL_NAMES = sorted(LAYOUT_BY_MODEL_NAME) + + +def iter_sequence_model_names() -> Iterable[str]: + for field in dataclasses.fields(FullTextModels): + if not isinstance(field.type, type) or not issubclass(field.type, Model): + continue + assert field.name.endswith('_model') + yield field.name[:-len('_model')] + + +def get_data_generator_by_model_name() -> Dict[str, ModelDataGenerator]: + fulltext_models = MockFullTextModels() + return { + model_name: fulltext_models.get_sequence_model_by_name( + model_name + ).get_data_generator( + document_features_context=DEFAULT_DOCUMENT_FEATURES_CONTEXT + ) + for model_name in iter_sequence_model_names() + } + + +DATA_GENERATOR_BY_MODEL_NAME = get_data_generator_by_model_name() + + +class TestGrobidColumnLayoutFile: + def test_should_cover_every_sequence_model(self): + assert sorted(iter_sequence_model_names()) == MODEL_NAMES + + @pytest.mark.parametrize('model_name', MODEL_NAMES) + def test_should_name_the_data_generator_the_model_uses(self, model_name: str): + data_generator = DATA_GENERATOR_BY_MODEL_NAME[model_name] + assert type(data_generator).__name__ == LAYOUT_BY_MODEL_NAME[model_name].generator + + @pytest.mark.parametrize('model_name', MODEL_NAMES) + def test_should_match_the_columns_the_data_generator_emits(self, model_name: str): + layout = LAYOUT_BY_MODEL_NAME[model_name] + data_generator = DATA_GENERATOR_BY_MODEL_NAME[model_name] + assert ( + list(data_generator.feature_names) + == layout.get_data_generator_column_names() + ) + + @pytest.mark.parametrize('model_name', MODEL_NAMES) + def test_should_start_the_columns_with_the_token(self, model_name: str): + assert LAYOUT_BY_MODEL_NAME[model_name].columns[0] == 'token_text' + + @pytest.mark.parametrize('model_name', MODEL_NAMES) + def test_should_not_record_the_placeholder_as_a_grobid_column(self, model_name: str): + layout = LAYOUT_BY_MODEL_NAME[model_name] + assert PLACEHOLDER_COLUMN_NAME not in layout.columns + assert PLACEHOLDER_COLUMN_NAME not in layout.extra_columns + + @pytest.mark.parametrize('model_name', MODEL_NAMES) + def test_should_record_a_reference_training_corpus(self, model_name: str): + assert LAYOUT_BY_MODEL_NAME[model_name].reference_training_corpus + + @pytest.mark.parametrize('model_name', MODEL_NAMES) + def test_should_emit_the_expected_training_data_column_count(self, model_name: str): + layout = LAYOUT_BY_MODEL_NAME[model_name] + assert ( + len(layout.get_training_data_column_names()) + 1 + == EXPECTED_TRAINING_DATA_COLUMN_COUNT_BY_MODEL_NAME[model_name] + ) + + @pytest.mark.parametrize('model_name', MODEL_NAMES) + def test_should_emit_the_expected_column_count_with_extra_columns(self, model_name: str): + layout = LAYOUT_BY_MODEL_NAME[model_name] + assert ( + len(layout.get_training_data_column_names(include_extra_columns=True)) + 1 + == EXPECTED_TRAINING_DATA_COLUMN_COUNT_WITH_EXTRA_COLUMNS_BY_MODEL_NAME[model_name] + ) + + def test_should_only_record_extra_columns_for_segmentation(self): + assert { + model_name + for model_name, layout in LAYOUT_BY_MODEL_NAME.items() + if layout.extra_columns + } == {'segmentation'} + + +class TestGrobidColumnLayout: + def test_should_reject_an_unknown_label_slot(self): + with pytest.raises(ValueError): + GrobidColumnLayout( + name='model1', + generator='DataGenerator1', + label_slot='other', + columns=('token_text', 'feature1') + ) + + def test_should_replace_the_placeholder_with_the_label_when_filled(self): + layout = GrobidColumnLayout( + name='model1', + generator='DataGenerator1', + label_slot=LabelSlot.FILLED, + columns=('token_text', 'feature1') + ) + assert layout.get_data_generator_column_names() == [ + 'token_text', 'feature1', PLACEHOLDER_COLUMN_NAME + ] + assert layout.get_training_data_column_names() == ['token_text', 'feature1'] + assert layout.get_training_data_feature_indices() == [0] + + def test_should_keep_the_placeholder_before_the_label_when_unfilled(self): + layout = GrobidColumnLayout( + name='model1', + generator='DataGenerator1', + label_slot=LabelSlot.UNFILLED, + columns=('token_text', 'feature1') + ) + assert layout.get_data_generator_column_names() == [ + 'token_text', 'feature1', PLACEHOLDER_COLUMN_NAME + ] + assert layout.get_training_data_column_names() == [ + 'token_text', 'feature1', PLACEHOLDER_COLUMN_NAME + ] + assert layout.get_training_data_feature_indices() == [0, 1] + + def test_should_have_no_placeholder_when_absent(self): + layout = GrobidColumnLayout( + name='model1', + generator='DataGenerator1', + label_slot=LabelSlot.ABSENT, + columns=('token_text', 'feature1') + ) + assert layout.get_data_generator_column_names() == ['token_text', 'feature1'] + assert layout.get_training_data_column_names() == ['token_text', 'feature1'] + assert layout.get_training_data_feature_indices() == [0] + + def test_should_drop_extra_columns_unless_asked_for(self): + layout = GrobidColumnLayout( + name='model1', + generator='DataGenerator1', + label_slot=LabelSlot.ABSENT, + columns=('token_text', 'feature1'), + extra_columns=('extra1',) + ) + assert layout.get_data_generator_column_names() == [ + 'token_text', 'feature1', 'extra1' + ] + assert layout.get_training_data_column_names() == ['token_text', 'feature1'] + assert layout.get_training_data_feature_indices() == [0] + assert layout.get_training_data_column_names(include_extra_columns=True) == [ + 'token_text', 'feature1', 'extra1' + ] + assert layout.get_training_data_feature_indices(include_extra_columns=True) == [0, 1] + + def test_should_keep_the_placeholder_last_with_extra_columns(self): + layout = GrobidColumnLayout( + name='model1', + generator='DataGenerator1', + label_slot=LabelSlot.UNFILLED, + columns=('token_text', 'feature1'), + extra_columns=('extra1',) + ) + assert layout.get_data_generator_column_names() == [ + 'token_text', 'feature1', 'extra1', PLACEHOLDER_COLUMN_NAME + ] + assert layout.get_training_data_feature_indices() == [0, 2] + assert layout.get_training_data_feature_indices(include_extra_columns=True) == [0, 1, 2] + + +class TestGetGrobidColumnLayoutForModelName: + def test_should_reject_an_unknown_model(self): + with pytest.raises(ValueError): + get_grobid_column_layout_for_model_name('model1') + + def test_should_share_one_layout_between_the_two_name_models(self): + assert ( + get_grobid_column_layout_for_model_name('name_header').columns + == get_grobid_column_layout_for_model_name('name_citation').columns + ) + + +class TestGetValidatedTrainingDataFeatureIndices: + @pytest.fixture(name='layout') + def _layout(self) -> GrobidColumnLayout: + return GrobidColumnLayout( + name='model1', + generator='DataGenerator1', + label_slot=LabelSlot.FILLED, + columns=('token_text', 'feature1') + ) + + def test_should_accept_the_recorded_columns(self, layout: GrobidColumnLayout): + assert get_validated_training_data_feature_indices( + layout, + feature_column_count=2, + data_generator_name='DataGenerator1', + data_generator_column_names=['token_text', 'feature1', PLACEHOLDER_COLUMN_NAME] + ) == [0] + + def test_should_reject_another_feature_column_count(self, layout: GrobidColumnLayout): + with pytest.raises(ValueError): + get_validated_training_data_feature_indices(layout, feature_column_count=3) + + def test_should_reject_another_data_generator(self, layout: GrobidColumnLayout): + with pytest.raises(ValueError): + get_validated_training_data_feature_indices( + layout, + feature_column_count=2, + data_generator_name='DataGenerator2' + ) + + def test_should_reject_renamed_columns(self, layout: GrobidColumnLayout): + with pytest.raises(ValueError): + get_validated_training_data_feature_indices( + layout, + feature_column_count=2, + data_generator_column_names=[ + 'token_text', 'feature2', PLACEHOLDER_COLUMN_NAME + ] + )