micm_nlp.training.trainers

Trainer subclass behaviour: optimizer groups, dataloaders, samplers.

custom_trainer_class_factory mixes CustomTrainerMixin into whichever HF Trainer trainer.cls names, so everything below applies to Trainer and Seq2SeqTrainer alike.

What the mixin changes:

• Optimizer construction — custom_training_args.optimizer_grouped_parameters gives a learning rate and weight decay to parameters whose names contain the listed substrings; unmatched parameters fall back to the global values. This is how prompt embeddings get a schedule of their own.

• Dataloaders — train, eval and test dataloaders can batch to a token budget via TokenBudgetBatchSampler instead of a fixed sample count, with the budget calibrated lazily and cached per stage.

• Column retention — _remove_unused_columns is overridden so the length column survives long enough for the samplers that read it.

• Constrained generation — a generation_whitelist injects ConstrainedPrefixLogitsProcessor into the generation call.

RandomTaskExclusionBatchSampler builds batches that hold out a random task, for custom_training_args.random_task_exclusion.

Attributes

logger

Classes

CustomTrainerMixin	A mixin class that adds custom functionality to either a Trainer or Seq2SeqTrainer instance.
RandomTaskExclusionBatchSampler	A custom BatchSampler that:

Functions

build_inference_dataloader_kwargs(→ dict)	Construct DataLoader kwargs for eval/test.
custom_trainer_class_factory(BaseTrainer)	Mix CustomTrainerMixin into whichever Trainer class is configured.

Module Contents

class micm_nlp.training.trainers.CustomTrainerMixin

A mixin class that adds custom functionality to either a Trainer or Seq2SeqTrainer instance. I think this is copied from transformers 4.39.1

create_optimizer()

Setup the optimizer.

We provide a reasonable default that works well. If you want to use something else, you can pass a tuple in the Trainer’s init through optimizers, or subclass and override this method in a subclass.

create_optimizer_and_scheduler(num_training_steps: int)

Build optimizer and scheduler, then inspect the weight decay once.

The inspection prints which parameters landed in which group – with optimizer_grouped_parameters in play, a prompt-embedding group silently not matching is otherwise invisible.

get_eval_dataloader(eval_dataset: torch.utils.data.Dataset | None = None) → torch.utils.data.DataLoader

Returns the evaluation [~torch.utils.data.DataLoader].

Subclass and override this method if you want to inject some custom behavior.

Parameters:

eval_dataset (torch.utils.data.Dataset, optional) – If provided, will override self.eval_dataset. If it is a [~datasets.Dataset], columns not accepted by the model.forward() method are automatically removed. It must implement __len__.

get_test_dataloader(test_dataset: torch.utils.data.Dataset) → torch.utils.data.DataLoader

Returns the test DataLoader. When test_max_tokens_per_batch is set, uses TokenBudgetBatchSampler; otherwise falls back to the legacy fixed-batch path with SequentialSampler / LengthGroupedSampler.

get_train_dataloader() → torch.utils.data.DataLoader

Returns the training [~torch.utils.data.DataLoader].

Will use no sampler if train_dataset does not implement __len__, a random sampler (adapted to distributed training if necessary) otherwise.

Subclass and override this method if you want to inject some custom behavior.

prediction_step(model: torch.nn.Module, inputs: dict[str, torch.Tensor | Any], prediction_loss_only: bool, ignore_keys: list[str] | None = None, **gen_kwargs) → tuple[float | None, torch.Tensor | None, torch.Tensor | None]

Run one prediction step, applying the generation whitelist if configured.

When custom_training_args.generation_whitelist is set, a ConstrainedPrefixLogitsProcessor is injected for this step, restricting generation to the allowed strings.

class micm_nlp.training.trainers.RandomTaskExclusionBatchSampler(dataset, batch_size: int, drop_last: bool = False)

Bases: torch.utils.data.BatchSampler

A custom BatchSampler that: - Randomly excludes tasks per batch - Samples each example at most once per epoch - Continues until all tasks are exhausted (even if only 1 or 2 tasks remain)

Parameters:

• dataset – dataset with a task-id column, used to group indices by task.

• batch_size – examples per batch.

• drop_last – drop a trailing partial batch.

batch_size

drop_last = False

num_batches

original_indices_by_task

task_ids = []

total_examples

micm_nlp.training.trainers.build_inference_dataloader_kwargs(*, dataset, args, data_collator, token_budget: int | None) → dict

Construct DataLoader kwargs for eval/test.

When token_budget is None, returns the legacy fixed-batch kwargs (caller still needs to add sampler). When token_budget is an int, returns kwargs using a TokenBudgetBatchSampler instead of batch_size + sampler (those keys are omitted because PyTorch rejects them alongside batch_sampler).

micm_nlp.training.trainers.custom_trainer_class_factory(BaseTrainer: transformers.Trainer | transformers.Seq2SeqTrainer)

Mix CustomTrainerMixin into whichever Trainer class is configured.

A factory rather than a fixed subclass because trainer.cls names the base at runtime – Trainer, Seq2SeqTrainer, or another – and the mixin has to sit in front of it in the MRO either way.

Parameters:

BaseTrainer – the HuggingFace trainer class to extend.

Returns:

a new class accepting the usual arguments plus custom_args.

micm_nlp.training.trainers.logger