From a651d1355cc659082e0a9c67c60adc58221f5099 Mon Sep 17 00:00:00 2001 From: Daniel Ecer Date: Fri, 26 Jun 2026 14:20:19 +0100 Subject: [PATCH] Add --models flag to restrict training data generation per model The CLI now accepts --models ... to limit which model generators run. When omitted, all models are generated. The Makefile default (TRAINING_DATA_MODELS) restricts to the five currently validated models (segmentation, header, affiliation-address, reference-segmenter, citation); override the variable to enable additional ones. --- Makefile | 3 + .../training/cli/generate_data.py | 81 ++++++++++++++++--- 2 files changed, 71 insertions(+), 13 deletions(-) diff --git a/Makefile b/Makefile index 39c714fb..55b325dd 100644 --- a/Makefile +++ b/Makefile @@ -62,6 +62,8 @@ TRAINING_DATA_NUM_WORKERS ?= 1 # Per-document timeout in seconds; 0 disables. Skips outlier PDFs (e.g. 73-page, 38 MB) # that cause the JATS aligner to run for many minutes. TRAINING_DATA_DOCUMENT_TIMEOUT ?= 120 +# Models to generate training data for (space-separated). Override to add more. +TRAINING_DATA_MODELS ?= segmentation header affiliation-address reference-segmenter citation # Source training data (PDF + JATS XML) downloaded from the HF dataset. # TRAINING_DATA_OUTPUT must point to a checkout of the output repo; create a @@ -308,6 +310,7 @@ dev-generate-training-data: --split $(SOURCE_TRAINING_SPLIT) \ --num-workers $(TRAINING_DATA_NUM_WORKERS) \ --document-timeout $(TRAINING_DATA_DOCUMENT_TIMEOUT) \ + --models $(TRAINING_DATA_MODELS) \ --debug \ $(ARGS) diff --git a/sciencebeam_parser/training/cli/generate_data.py b/sciencebeam_parser/training/cli/generate_data.py index 01cbfebb..447819f1 100644 --- a/sciencebeam_parser/training/cli/generate_data.py +++ b/sciencebeam_parser/training/cli/generate_data.py @@ -145,6 +145,19 @@ def parse_args(argv: Optional[List[str]] = None) -> argparse.Namespace: 'Single-worker mode uses SIGALRM; multi-worker mode uses future timeout.' ) ) + parser.add_argument( + '--models', + type=str, + nargs='+', + default=None, + metavar='MODEL', + help=( + 'Models to generate training data for. ' + 'If omitted, all models are generated. ' + 'Valid names: segmentation, header, affiliation-address, name-header, ' + 'fulltext, figure, table, reference-segmenter, citation, name-citation.' + ) + ) return parser.parse_args(argv) @@ -404,6 +417,11 @@ def _apply_jats_labels_to_model_data_list( class AbstractModelTrainingDataGenerator(ABC): + @property + @abstractmethod + def model_name(self) -> str: + pass + def get_pre_file_path_suffix(self) -> str: return '' @@ -550,6 +568,8 @@ def iter_model_data_list( class SegmentationModelTrainingDataGenerator(AbstractDocumentModelTrainingDataGenerator): + model_name = 'segmentation' + def get_main_model(self, document_context: TrainingDataDocumentContext) -> Model: return document_context.fulltext_models.segmentation_model @@ -571,6 +591,8 @@ def fn( class HeaderModelTrainingDataGenerator(AbstractDocumentModelTrainingDataGenerator): + model_name = 'header' + def get_main_model(self, document_context: TrainingDataDocumentContext) -> Model: return document_context.fulltext_models.header_model @@ -638,6 +660,8 @@ def iter_model_layout_documents( class AffiliationAddressModelTrainingDataGenerator(AbstractDocumentModelTrainingDataGenerator): + model_name = 'affiliation-address' + def get_main_model(self, document_context: TrainingDataDocumentContext) -> Model: return document_context.fulltext_models.affiliation_address_model @@ -695,6 +719,8 @@ def iter_model_layout_documents( class NameHeaderModelTrainingDataGenerator(AbstractDocumentModelTrainingDataGenerator): + model_name = 'name-header' + def get_main_model(self, document_context: TrainingDataDocumentContext) -> Model: return document_context.fulltext_models.name_header_model @@ -749,6 +775,8 @@ def iter_model_layout_documents( class NameCitationModelTrainingDataGenerator(AbstractDocumentModelTrainingDataGenerator): + model_name = 'name-citation' + def get_main_model(self, document_context: TrainingDataDocumentContext) -> Model: return document_context.fulltext_models.name_citation_model @@ -830,6 +858,8 @@ def iter_model_layout_documents( class FullTextModelTrainingDataGenerator(AbstractDocumentModelTrainingDataGenerator): + model_name = 'fulltext' + def get_main_model(self, document_context: TrainingDataDocumentContext) -> Model: return document_context.fulltext_models.fulltext_model @@ -864,6 +894,8 @@ def iter_model_layout_documents( class FigureModelTrainingDataGenerator(AbstractDocumentModelTrainingDataGenerator): + model_name = 'figure' + def get_main_model(self, document_context: TrainingDataDocumentContext) -> Model: return document_context.fulltext_models.figure_model @@ -905,6 +937,8 @@ def iter_model_layout_documents( class TableModelTrainingDataGenerator(AbstractDocumentModelTrainingDataGenerator): + model_name = 'table' + def get_main_model(self, document_context: TrainingDataDocumentContext) -> Model: return document_context.fulltext_models.table_model @@ -946,6 +980,8 @@ def iter_model_layout_documents( class ReferenceSegmenterModelTrainingDataGenerator(AbstractDocumentModelTrainingDataGenerator): + model_name = 'reference-segmenter' + def get_main_model(self, document_context: TrainingDataDocumentContext) -> Model: return document_context.fulltext_models.reference_segmenter_model @@ -967,6 +1003,8 @@ def iter_model_layout_documents( class CitationModelTrainingDataGenerator(AbstractDocumentModelTrainingDataGenerator): + model_name = 'citation' + def get_main_model(self, document_context: TrainingDataDocumentContext) -> Model: return document_context.fulltext_models.citation_model @@ -1023,6 +1061,28 @@ def iter_model_layout_documents( ] +def _select_generators( + enabled_models: Optional[frozenset], +) -> List['AbstractModelTrainingDataGenerator']: + all_gens: List[AbstractModelTrainingDataGenerator] = [ + SegmentationModelTrainingDataGenerator(), + HeaderModelTrainingDataGenerator(), + AffiliationAddressModelTrainingDataGenerator(), + NameHeaderModelTrainingDataGenerator(), + FullTextModelTrainingDataGenerator(), + FigureModelTrainingDataGenerator(), + TableModelTrainingDataGenerator(), + ReferenceSegmenterModelTrainingDataGenerator(), + CitationModelTrainingDataGenerator(), + NameCitationModelTrainingDataGenerator(), + ] + if enabled_models is None: + return all_gens + selected = [g for g in all_gens if g.model_name in enabled_models] + LOGGER.info('enabled models: %s', sorted(enabled_models)) + return selected + + def _build_jats_annotations( layout_document: LayoutDocument, jats_xml_filename: str, @@ -1049,6 +1109,7 @@ def generate_training_data_for_layout_document( use_directory_structure: bool, gzip_enabled: bool = False, jats_xml_filename: Optional[str] = None, + enabled_models: Optional[frozenset] = None, ): model_result_cache = ModelResultCache() jats_annotated: Optional[JatsAnnotatedLayoutDocument] = None @@ -1074,19 +1135,7 @@ def generate_training_data_for_layout_document( jats_annotated_document=jats_annotated, jats_segmentation_labels=jats_seg_labels, ) - training_data_generators = [ - SegmentationModelTrainingDataGenerator(), - HeaderModelTrainingDataGenerator(), - AffiliationAddressModelTrainingDataGenerator(), - NameHeaderModelTrainingDataGenerator(), - FullTextModelTrainingDataGenerator(), - FigureModelTrainingDataGenerator(), - TableModelTrainingDataGenerator(), - ReferenceSegmenterModelTrainingDataGenerator(), - CitationModelTrainingDataGenerator(), - NameCitationModelTrainingDataGenerator() - ] - for training_data_generator in training_data_generators: + 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 @@ -1133,6 +1182,7 @@ def generate_training_data_for_source_filename( use_directory_structure: bool, gzip_enabled: bool, xml_file_list: Optional[Sequence[str]] = None, + enabled_models: Optional[frozenset] = None, ): LOGGER.debug('use_model: %r', use_model) layout_document = get_layout_document_for_source_filename( @@ -1158,6 +1208,7 @@ def generate_training_data_for_source_filename( use_directory_structure=use_directory_structure, gzip_enabled=gzip_enabled, jats_xml_filename=jats_xml_filename, + enabled_models=enabled_models, ) @@ -1228,6 +1279,7 @@ def _worker_process(kwargs: dict) -> bool: use_directory_structure=kwargs['use_directory_structure'], gzip_enabled=kwargs['gzip_enabled'], xml_file_list=kwargs['xml_file_list'], + enabled_models=kwargs['enabled_models'], ) return True except Exception: # pylint: disable=broad-except @@ -1260,6 +1312,7 @@ def _run_serial( 'use_directory_structure': args.use_directory_structure, 'gzip_enabled': args.gzip, 'xml_file_list': xml_file_list, + 'enabled_models': args.enabled_models, } if document_timeout == 0: @@ -1320,6 +1373,7 @@ def _run_parallel_workers( 'use_directory_structure': args.use_directory_structure, 'gzip_enabled': args.gzip, 'xml_file_list': xml_file_list, + 'enabled_models': args.enabled_models, } # pylint: disable-next=consider-using-with pool = multiprocessing.Pool(num_workers, initializer=_worker_init) @@ -1360,6 +1414,7 @@ def run(args: argparse.Namespace): if args.source_xml_path: xml_file_list = list(glob(args.source_xml_path)) LOGGER.info('JATS XML files: %d', len(xml_file_list)) + args.enabled_models = frozenset(args.models) if args.models else None # Note: creating the directory may not be necessary, but provides early feedback makedirs(output_path, exist_ok=True) total = len(source_file_list)