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

对提供的 dataset 执行贪心样本打包。这作为训练开始前的一个单独预处理步骤完成。打包后的样本在 dataloader 中使用 Sampler 在此类外部进行打乱。当前，这仅支持内存中的 map-style 数据集。

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

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

训练期间，返回 self.packs[idx] 作为输入、标签、注意力掩码和位置 ID。注意力掩码是下三角块状掩码，用于阻止样本在包内相互注意。位置 ID 指示包内每个标记相对于其样本的位置。这些都填充到最大序列长度，因此不需要 batch-wise 的 collator。

一个打包样本由多个较小序列长度的样本在 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，我们只是从上一个样本正常继续递增。

参数：

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

PackedDataset¶

文档

教程

资源