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)[源代码]¶

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

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

初始化的通用流程是：加载分词样本 -> 添加到缓冲区 -> 当缓冲区足够长时，添加到 self.packs。

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

一个打包的样本由 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],
    ...,
]

在 pad token 的掩码中使用单位矩阵而不是因果掩码。对于 pad token 的位置 ID，我们只需从上一个样本正常递增。

参数:

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

PackedDataset¶

文档

教程

资源