Collate

bert_padding_collate_fn(batch, padding_value, min_length=None, max_length=None)

BERT 数据加载器的 Padding collate 函数。

参数

名称	类型	描述	默认值
batch	list	样本列表。	必需
padding_value	int	Tokenizer 的 pad token ID。	必需
min_length	int | None	输出批次的最小长度；张量将被填充到此长度。如果未提供，则不会添加超出 max_length 的额外填充。	None
max_length	int | None	序列的最大长度。如果未提供，张量将被填充到批次中最长序列的长度。	None

源代码位于 bionemo/llm/data/collate.py 中

def bert_padding_collate_fn(
    batch: Sequence[types.BertSample],
    padding_value: int,
    min_length: int | None = None,
    max_length: int | None = None,
) -> types.BertSample:
    """Padding collate function for BERT dataloaders.

    Args:
        batch (list): List of samples.
        padding_value (int, optional): The tokenizer's pad token ID.
        min_length: Minimum length of the output batch; tensors will be padded to this length. If not
            provided, no extra padding beyond the max_length will be added.
        max_length: Maximum length of the sequence. If not provided, tensors will be padded to the
            longest sequence in the batch.
    """
    padding_values = {
        "text": padding_value,
        "types": 0,
        "attention_mask": False,
        "labels": MLM_LOSS_IGNORE_INDEX,  # This should match the masked value used in the MLM loss mask.
        "loss_mask": False,
        "is_random": 0,
    }
    return padding_collate_fn(
        batch=batch,  # type: ignore[assignment]
        padding_values=padding_values,
        min_length=min_length,
        max_length=max_length,
    )

padding_collate_fn(batch, padding_values, min_length=None, max_length=None)

带有 padding 的 Collate 函数。

参数

名称	类型	描述	默认值
batch	Sequence[_T]	样本列表，每个样本都是张量字典。	必需
padding_values	dict[str, int]	每个张量键的 padding 值字典。	必需
min_length	int | None	输出批次的最小长度；张量将被填充到此长度。如果未提供，则不会添加超出 max_length 的额外填充。	None
max_length	int | None	序列的最大长度。如果未提供，张量将被填充到批次中最长序列的长度。	None

返回

类型	描述
_T	具有相同字典输入结构的 collate 批次。

源代码位于 bionemo/llm/data/collate.py 中

def padding_collate_fn(
    batch: Sequence[_T],
    padding_values: dict[str, int],
    min_length: int | None = None,
    max_length: int | None = None,
) -> _T:
    """Collate function with padding.

    Args:
        batch: List of samples, each of which is a dictionary of tensors.
        padding_values: A dictionary of padding values for each tensor key.
        min_length: Minimum length of the output batch; tensors will be padded to this length. If not
            provided, no extra padding beyond the max_length will be added.
        max_length: Maximum length of the sequence. If not provided, tensors will be padded to the
            longest sequence in the batch.

    Returns:
        A collated batch with the same dictionary input structure.
    """
    global _warned_once
    keys: set[str] | None = None

    if len(batch) == 0:  # empty batches passed through in DDP inference
        return {}

    for entry in batch:
        # First check that we have sane batches where keys align with each other.
        if keys is None:
            keys = set(entry.keys())
        else:
            if set(entry.keys()) != keys:
                raise ValueError(f"All keys in inputs must match each other. Got: {[sorted(e.keys()) for e in batch]}")
        if entry.keys() != padding_values.keys():
            if not _warned_once:
                extra_keys = {k for k in entry.keys() if k not in padding_values}
                missing_keys = {k for k in padding_values.keys() if k not in entry}
                logger.warning(
                    f"Extra keys in batch that will not be padded: {extra_keys}. Missing keys in batch: {missing_keys}"
                )
                _warned_once = True

    def _pad(tensors, padding_value):
        if max_length is not None:
            tensors = [t[:max_length] for t in tensors]
        batched_tensors = torch.nn.utils.rnn.pad_sequence(tensors, batch_first=True, padding_value=padding_value)
        if min_length is None:
            return batched_tensors
        return torch.nn.functional.pad(batched_tensors, (0, min_length - batched_tensors.size(1)), value=padding_value)

    return {
        k: _pad([s[k] for s in batch], padding_values[k])
        if k in padding_values
        else torch.stack([s[k] for s in batch])
        for k in batch[0].keys()
    }  # type: ignore[return-value]