micm_nlp.training.data_collators

Data collators beyond the HuggingFace defaults.

A collator is selected by name from data_collator.cls, resolved against transformers and this module, with data_collator.args handed to its constructor.

  • DataCollatorForPLMWithPadding — permutation language modelling with padding.

  • DataCollatorForSeq2SeqWithShiftLabels — seq2seq collation that shifts labels by a configurable offset.

  • CustomDataCollatorForSeq2Seq_2 and CustomDataCollatorWithPadding — variants carrying the padding behaviour this toolkit needs.

  • DataCollatorTaskIDDecorator — wraps another collator to attach a task id.

DataCollatorTaskIDDecorator has no consumers yet; its leftover print/exit debug body was removed in 0.2.0, so it now does what its name says, but it has not been exercised in a real run.

Classes

CustomDataCollatorForSeq2Seq_2

Data collator that will dynamically pad the inputs received, as well as the labels.

CustomDataCollatorWithPadding

Data collator that will dynamically pad the inputs received.

DataCollatorForPLMWithPadding

Permutation language modelling (XLNet) with padding applied first.

DataCollatorForSeq2SeqWithShiftLabels

Seq2seq collation that pads the labels to account for virtual tokens.

DataCollatorTaskIDDecorator

Decorator for adding task_ids to any data collator.

Module Contents

class micm_nlp.training.data_collators.CustomDataCollatorForSeq2Seq_2

Bases: transformers.DataCollatorForSeq2Seq

Data collator that will dynamically pad the inputs received, as well as the labels.

Parameters:
  • tokenizer ([PreTrainedTokenizer] or [PreTrainedTokenizerFast]) – The tokenizer used for encoding the data.

  • model ([PreTrainedModel], optional) –

    The model that is being trained. If set and has the prepare_decoder_input_ids_from_labels, use it to prepare the decoder_input_ids

    This is useful when using label_smoothing to avoid calculating loss twice.

  • padding (bool, str or [~utils.PaddingStrategy], optional, defaults to True) –

    Select a strategy to pad the returned sequences (according to the model’s padding side and padding index) among:

    • True or ‘longest’ (default): Pad to the longest sequence in the batch (or no padding if only a single sequence is provided).

    • ’max_length’: Pad to a maximum length specified with the argument max_length or to the maximum acceptable input length for the model if that argument is not provided.

    • False or ‘do_not_pad’: No padding (i.e., can output a batch with sequences of different lengths).

  • max_length (int, optional) – Maximum length of the returned list and optionally padding length (see above).

  • pad_to_multiple_of (int, optional) –

    If set will pad the sequence to a multiple of the provided value.

    This is especially useful to enable the use of Tensor Cores on NVIDIA hardware with compute capability >= 7.0 (Volta).

  • label_pad_token_id (int, optional, defaults to -100) – The id to use when padding the labels (-100 will be automatically ignored by PyTorch loss functions).

  • return_tensors (str, optional, defaults to “pt”) – The type of Tensor to return. Allowable values are “np”, “pt” and “tf”.

label_pad_token_id: int = -100
max_length: int | None = None
model: Any | None = None
pad_to_multiple_of: int | None = None
padding: bool | str | transformers.tokenization_utils_base.PaddingStrategy = True
return_tensors: str = 'pt'
tokenizer: transformers.tokenization_utils_base.PreTrainedTokenizerBase
class micm_nlp.training.data_collators.CustomDataCollatorWithPadding

Bases: transformers.DataCollatorWithPadding

Data collator that will dynamically pad the inputs received.

Parameters:
  • tokenizer ([PreTrainedTokenizer] or [PreTrainedTokenizerFast]) – The tokenizer used for encoding the data.

  • padding (bool, str or [~utils.PaddingStrategy], optional, defaults to True) –

    Select a strategy to pad the returned sequences (according to the model’s padding side and padding index) among:

    • True or ‘longest’ (default): Pad to the longest sequence in the batch (or no padding if only a single sequence is provided).

    • ’max_length’: Pad to a maximum length specified with the argument max_length or to the maximum acceptable input length for the model if that argument is not provided.

    • False or ‘do_not_pad’: No padding (i.e., can output a batch with sequences of different lengths).

  • max_length (int, optional) – Maximum length of the returned list and optionally padding length (see above).

  • pad_to_multiple_of (int, optional) –

    If set will pad the sequence to a multiple of the provided value.

    This is especially useful to enable the use of Tensor Cores on NVIDIA hardware with compute capability >= 7.0 (Volta).

  • return_tensors (str, optional, defaults to “pt”) – The type of Tensor to return. Allowable values are “np”, “pt” and “tf”.

max_length: int | None = None
pad_to_multiple_of: int | None = None
padding: bool | str | transformers.tokenization_utils_base.PaddingStrategy = True
return_tensors: str = 'pt'
tokenizer: transformers.tokenization_utils_base.PreTrainedTokenizerBase
class micm_nlp.training.data_collators.DataCollatorForPLMWithPadding(tokenizer: transformers.XLNetTokenizer, max_length: int = 128, padding: str = 'max_length', plm_probability: float = 1 / 6, max_span_length: int = 5, return_tensors: str = 'pt')

Bases: transformers.DataCollatorForPermutationLanguageModeling

Permutation language modelling (XLNet) with padding applied first.

HuggingFace’s DataCollatorForPermutationLanguageModeling assumes every sequence in the batch is already the same length – it builds the permutation mask and target mapping from a fixed shape. This pads the batch to max_length (or to the longest row), then hands the padded batch to the parent, so variable-length examples can be used.

Padding is forced to the right regardless of the tokenizer’s setting, because the permutation mask is built from the left.

Parameters:
  • tokenizer – an XLNet tokenizer.

  • max_length – target length when padding='max_length'.

  • padding'max_length', True (longest in batch) or False.

  • plm_probability – fraction of tokens to mask, as the parent defines it.

  • max_span_length – longest span of masked tokens.

  • return_tensors – framework for the returned batch.

torch_call(features)

Pad the batch, then build the permutation mask and target mapping.

max_length = 128
padding = 'max_length'
return_tensors = 'pt'
tokenizer
class micm_nlp.training.data_collators.DataCollatorForSeq2SeqWithShiftLabels(*args, shift_labels_by=None, **kwargs)

Bases: transformers.DataCollatorForSeq2Seq

Seq2seq collation that pads the labels to account for virtual tokens.

Prompt learning prepends N virtual tokens to the decoder input, which makes the labels one position short at every step. This left-pads them with -100 – ignored by the loss – so labels line up with the shifted inputs.

shift_labels_by should equal the number of virtual tokens; None leaves the labels alone, making this behave as the stock collator.

Parameters:
  • args – forwarded to DataCollatorForSeq2Seq.

  • shift_labels_by – positions to left-pad the labels by, normally the virtual-token count.

  • kwargs – forwarded to DataCollatorForSeq2Seq.

shift_labels_by = None
class micm_nlp.training.data_collators.DataCollatorTaskIDDecorator(base_collator, task_id)

Decorator for adding task_ids to any data collator.

Parameters:
  • base_collator – the collator to wrap; any collator will do.

  • task_id – task id to attach to every batch this produces.

base_collator
task_id