PackedDataset¶

class torchtune.datasets.PackedDataset(ds: Dataset, *, max_seq_len: int, padding_idx: int = 0, max_packs: Optional[int] = None, split_across_pack: bool = False)[source]¶

对提供的数据集执行贪婪样本打包。这是在训练开始前作为单个预处理步骤完成的。洗牌是在此类之外使用Sampler作为数据加载器的一部分对打包样本进行的。目前，这仅支持内存中的映射式数据集。

该类在初始化时加载、分词并打包示例 - 在训练期间不进行分词。

初始化的一般流程是：加载标记化样本 -> 添加到缓冲区 -> 当缓冲区足够长时，添加到 self.packs。

在训练期间，返回 self.packs[idx] 作为输入、标签、注意力掩码和位置ID。注意力掩码是一个下三角块掩码，用于防止样本在包内进行交叉注意力。位置ID表示每个标记相对于其在包内样本的位置。这些都填充到最大序列长度，因此不需要批处理整理器。

一个打包的样本由多个较小的序列长度样本挤在一起组成，这些样本在max_seq_len内。例如，如果max_seq_len是6，并且有不同长度的样本：

tokens = [
    [S1, S1, S1, S2, S2, pad],
    [S3, S3, S4, S4, pad, pad],
    ...,
]

为了防止交叉污染，示例中的第一个包将返回以下掩码：

mask = [
    [1, 0, 0, 0, 0, 0],
    [1, 1, 0, 0, 0, 0],
    [1, 1, 1, 0, 0, 0],
    [0, 0, 0, 1, 0, 0],
    [0, 0, 0, 1, 1, 0],
    [0, 0, 0, 0, 0, 1],
]

位置ID将是：

input_pos = [
    [0, 1, 2, 0, 1, 2],
    [0, 1, 0, 1, 2, 3],
    ...,
]

单位矩阵用于填充标记的掩码，而不是因果掩码。对于填充标记的位置ID，我们只需从上一个样本正常继续递增。

Parameters:

ds (Dataset) – 用于采样的数据集。这应该返回一个包含字段“tokens”和“labels”的字典，其中包含标记化和标签样本。
max_seq_len (int) – 最大打包的标记数量
padding_idx (int) – 分词器的填充索引。默认值为0。
max_packs (可选[int]) – 最大包数。默认为 None，这将创建尽可能多的包。
split_across_pack (bool) – 如果一个包中的最后一个样本不适合max_seq_len，则将样本分割到下一个包中，或将其完全移动到下一个包的开头。对于预训练，通常将其设置为True以进行一般文本完成。对于微调，通常将其设置为False以避免在指令调优中截断句子。默认值为False。