Skip to content

Add --models flag to restrict training data generation per model #679

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
de-code merged 1 commit into from
Jun 26, 2026
Merged

Add --models flag to restrict training data generation per model #679

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions Makefile
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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)

Expand Down
81 changes: 68 additions & 13 deletions sciencebeam_parser/training/cli/generate_data.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)


Expand Down Expand Up @@ -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 ''

Expand Down Expand Up @@ -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

Expand All @@ -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

Expand Down Expand Up @@ -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

Expand Down Expand Up @@ -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

Expand Down Expand Up @@ -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

Expand Down Expand Up @@ -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

Expand Down Expand Up @@ -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

Expand Down Expand Up @@ -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

Expand Down Expand Up @@ -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

Expand All @@ -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

Expand Down Expand Up @@ -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,
Expand All @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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(
Expand All @@ -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,
)


Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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)
Expand Down
Loading