micm_nlp.training.runner

TRAINER — assemble the HuggingFace Trainer and drive the run.

run() is the entry point: evaluate before training if asked, train, evaluate after, optionally run the test split, and return the test output. Everything else in the class is assembly — training arguments, the data collator, the metrics callables, callbacks, W&B initialisation, and evaluation-step arithmetic.

Two details worth knowing:

Selection and stopping are separate. training_args.metric_for_best_model picks the checkpoint; custom_training_args.early_stopping_metric decides what stops the run. They can differ — useful when the evaluation loss is unstable but accuracy is the thing you care about.

Prediction order is not always dataset order. Token-budget batching length-sorts globally, so _aligned_ds_split permutes the split to the sampler’s emit order before predictions are zipped against it. Samplers opt in simply by exposing an order property.

Classes

TRAINER

Builds the HuggingFace Trainer from the config, and runs it.

Module Contents

class micm_nlp.training.runner.TRAINER(model, dataset, tokenizer=None)

Builds the HuggingFace Trainer from the config, and runs it.

Everything the Trainer needs is assembled here from YAML: the training arguments class, the collator and the Trainer subclass are all resolved by name against the *_SOURCE_MODULES lists, which is why a run can select a HuggingFace class or one of this package’s without any code change.

Construction wires the trainer; run() executes the phases the config asks for – zero-shot test, evaluation before training, training, evaluation after, final test – in that order.

Assemble the Trainer and print what was built.

Parameters:

  • model – a MODEL; its config drives everything here.

  • dataset – a DATASET.

  • tokenizer – tokenizer to use; defaults to the dataset’s own.

print_batch_examples(split, batches=2, samples_per_batch=2)

Print a few decoded examples as the model will actually receive them.

Goes through the real dataloader, so what is printed reflects collation, padding and any special tokens – the fastest way to catch a prompt template or label alignment that is wrong. Returns quietly if the split is absent.

Parameters:

  • split – which split to sample, a DsSplitSE.

  • batches – how many batches to show.

  • samples_per_batch – how many rows from each.

print_details()

Print the assembled model, dataset and trainer setup.

The record of what a run actually built – resolved classes, paths, sizes – as opposed to what the YAML asked for.

run()

Run the phases the config selects, in order.

Zero-shot test, evaluation before training, training, evaluation after training, then the final test. test.zero_shot_only skips training entirely, which is how a zero-shot baseline row is produced; each phase is otherwise gated by its own flag in eval / test.

Returns:

the test output – both the full-shot and zero-shot results when both ran.

DATA_COLLATOR_SOURCE_MODULES: ClassVar[list[str]] = ['transformers', 'micm_nlp.training.data_collators']
TRAINER_SOURCE_MODULES: ClassVar[list[str]] = ['transformers', 'micm_nlp.training.trainers']
TRAINING_ARGS_SOURCE_MODULES: ClassVar[list[str]] = ['transformers']