integrations.kd.collator

integrations.kd.collator

DataCollator for axolotl to handle KD fields without using -inf for padding, and with a teacher_mask to identify padded positions.

Classes

Name	Description
DataCollatorForKD	Data collator for KD, including handling KD-specific fields.
KDBatchSamplerDataCollatorForSeq2Seq	Collator for multipack (batch of sub-batches) specifically for KD.

DataCollatorForKD

integrations.kd.collator.DataCollatorForKD(*args, **kwargs)

Data collator for KD, including handling KD-specific fields.

This version avoids using -inf and instead uses a large negative value for padding target_logprobs. It also creates a teacher_mask to indicate which entries are valid.

KDBatchSamplerDataCollatorForSeq2Seq

integrations.kd.collator.KDBatchSamplerDataCollatorForSeq2Seq(*args, **kwargs)

Collator for multipack (batch of sub-batches) specifically for KD. Adapts DataCollatorForKD so it can pack multiple sequences in a single batch item.