LLM - DataCollatorForLanguageModeling 样本生成 by transformers

一.引言

二.生成样本 By API

1.样本处理样式

2.DataCollatorForLanguageModeling

2.1 样本准备

2.2 API 生成

三.生成样本 By DIY

1.样本准备

2.data_colloator 实现

3.使用自定义 data_colloator

四.总结

一.引言

前面我们讲了 Baichuan7B 的 lora 微调步骤，我们在 QA 基础上构建了样本集并训练，但是细心的同学肯定发现我们原始样本中只给出了 input_ids 但是没有给出 labels，本文我们简单看下 data_collator 如何生成样本 label 并自定义实现一个简单的 data_collator。

二.生成样本 By API

1.样本处理样式

def preprocess(tokenizer, config, example, max_seq_length, prompt_key, target_key):
    prompt = example[prompt_key]
    target = example[target_key]
    prompt_ids = tokenizer.encode(prompt, max_length=max_seq_length, truncation=True)
    target_ids = tokenizer.encode(target, max_length=max_seq_length, truncation=True, add_special_tokens=False)
    # 最终还是将 instruction 的输入输出都拼在一起，使用经典的 causal-LM 的 next word prediction 方式来训练
    input_ids = prompt_ids + target_ids + [config.eos_token_id] # EOS 用于标识句子结束
    return {"input_ids": input_ids, "seq_len": len(prompt_ids)}

首先加载 tokenizer 对 Q 和 A 分别 token 获取对应 Q_ids 和 A_ids：

Q: 请计算：39 * 0 = 什么？
A: 这是简单的乘法运算，39乘以0得到的是0
TokenQ: [31106, 15875, 77, 57, 53, 31159, 53, 60, 58, 31135, 6787, 6775, 81]
TokenA: [31106, 3908, 14313, 31640, 31257, 31481, 31742, 72, 57, 53, 31640, 31187, 53, 60, 58, 9323, 31178, 52, 79, 54, 59, 56]

获取 ids 后，会直接将 Q_ids A_ids 连接并在尾部增加 eos_token_id 标识当前句子结束，这里 eos_token_id = 2，合并后的 input_ids 如下：

input_ids = [31106, 15875, 77, 57, 53, 31159, 53, 60, 58, 31135, 6787, 6775, 81, 31106, 3908, 14313, 31640, 31257, 31481,
           31742, 72, 57, 53, 31640, 31187, 53, 60, 58, 9323, 31178, 52, 79, 54, 59, 56, 2]

除此之外，json 里还用 seq_len 记录了 prompt_ids 即 Q_ids 的长度：

json = {"input_ids": input_ids, "seq_len": 13}

2.DataCollatorForLanguageModeling

data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer,mlm=False)

    # 初始化 trainer, 此处报错: NotImplementedError: Cannot copy out of meta tensor; no data!
    trainer = ModifiedTrainer(
        model=model,
        train_dataset=dataset,
        args=training_args,
        callbacks=[TensorBoardCallback(writer)], # 回调将训练情况写入 tensorboard
        data_collator=data_collator,
    )
    trainer.train()

上文中我们把 json 样本生成的 Dataset 和 DataCollatorForLanguageModeling 直接传给了 Trainer 训练，但是样本中并没有 labels，于是就在想没有 labels 怎么梯度回传，下面用 DataCollatorForLanguageModeling API 看下 data_collator 使用后生成的数据样式。

2.1 样本准备

这里我直接将第一步 tokenizer 后的样本手动生成两个 json 测试后续流程：

# p: [31106, 15875, 77, 57, 53, 31159, 53, 60, 58, 31135, 6787, 6775, 81]
# t: [31106, 3908, 14313, 31640, 31257, 31481, 31742, 72, 57, 53, 31640, 31187, 53, 60, 58, 9323, 31178, 52, 79, 54, 59, 56]
sample1 = [31106, 15875, 77, 57, 53, 31159, 53, 60, 58, 31135, 6787, 6775, 81, 31106, 3908, 14313, 31640, 31257, 31481,
           31742, 72, 57, 53, 31640, 31187, 53, 60, 58, 9323, 31178, 52, 79, 54, 59, 56, 2]

# p: [31106, 33370, 5629, 16927, 54, 56, 31179, 11002, 72, 31526, 31342, 8835, 31221, 31423, 7156, 55, 31396, 31222, 33370, 31604, 72, 31415, 31396, 31222, 33370, 1231, 31221, 31195, 6851, 11002, 75]
# t: [31106, 33370, 5629, 16927, 54, 56, 31179, 11002, 72, 8835, 31221, 31423, 55, 31396, 31222, 33370, 31604, 72, 2926, 31415, 31396, 31222, 33370, 1231, 31221, 11254, 11002, 31380, 1522, 31482, 11002, 31380, 31640, 31187, 31222, 33370, 31135, 31396, 31380, 73, 5, 54, 56, 34399, 55, 31191, 60, 5, 31415, 31396, 31222, 33370, 1231, 31221, 31195, 60, 31179, 11002, 73, 2122, 72, 6787, 31161, 31415, 31396, 31222, 33370, 1231, 31221, 31195, 60, 31179, 11002, 73]
sample2 = [31106, 33370, 5629, 16927, 54, 56, 31179, 11002, 72, 31526, 31342, 8835, 31221, 31423, 7156, 55, 31396,
           31222, 33370, 31604, 72, 31415, 31396, 31222, 33370, 1231, 31221, 31195, 6851, 11002, 75, 31106, 33370, 5629,
           16927, 54, 56, 31179, 11002, 72, 8835, 31221, 31423, 55, 31396, 31222, 33370, 31604, 72, 2926, 31415, 31396,
           31222, 33370, 1231, 31221, 11254, 11002, 31380, 1522, 31482, 11002, 31380, 31640, 31187, 31222, 33370, 31135,
           31396, 31380, 73, 5, 54, 56, 34399, 55, 31191, 60, 5, 31415, 31396, 31222, 33370, 1231, 31221, 31195, 60,
           31179, 11002, 73, 2122, 72, 6787, 31161, 31415, 31396, 31222, 33370, 1231, 31221, 31195, 60, 31179, 11002,
           73, 2]

json1 = {"input_ids": sample1, "seq_len": 13}
json2 = {"input_ids": sample2, "seq_len": 31}

features = [json1, json2]

每个 json 的 input_ids 都遵循 Q_ids + A_ids + [SEP] 即 Prompt_ids + Target_ids + [SEP] 的规则生成。

2.2 API 生成

from transformers import AutoTokenizer
from transformers import DataCollatorForLanguageModeling

model_checkpoint = "/model/baichuan-7B"
tokenizer = AutoTokenizer.from_pretrained(model_checkpoint, trust_remote_code=True)
tokenizer.pad_token = tokenizer.unk_token
data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer,mlm=False)

features = [json1, json2]
batch = data_collator(features)

print(batch)

调用 API 生成一个 Batch 对应的大 json 共包含 4 个 key：

• attention_mask

对于基于 mask 的语言任务例如 Bert，data_collator 可以帮助生成掩码标签。对于 sample1 其长度较短，所以后面的 PAD_TOKEN 部分的 mask 均为 0，而 sample2 的长度为所有样本中最长的且小于 max_length，所以其 mask 均为 1。

'attention_mask': tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
		 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 
		 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 
		 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
		 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
         1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
         1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
         1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
         1, 1, 1, 1, 1, 1, 1, 1, 1, 1]])

• input_ids

与上面同理，data_collator 会就算当前 batch 样本中的 max_length，对于不足 max_length 的样本进行补齐的 padding 操作，所以 sample1 的样本补了很多 padding，sample2 正常。

'input_ids': tensor([[31106, 15875,    77,    57,    53, 31159,    53,    60,    58, 31135,
          6787,  6775,    81, 31106,  3908, 14313, 31640, 31257, 31481, 31742,
            72,    57,    53, 31640, 31187,    53,    60,    58,  9323, 31178,
            52,    79,    54,    59,    56,     2,     0,     0,     0,     0,
             0,     0,     0,     0,     0,     0,     0,     0,     0,     0,
             0,     0,     0,     0,     0,     0,     0,     0,     0,     0,
             0,     0,     0,     0,     0,     0,     0,     0,     0,     0,
             0,     0,     0,     0,     0,     0,     0,     0,     0,     0,
             0,     0,     0,     0,     0,     0,     0,     0,     0,     0,
             0,     0,     0,     0,     0,     0,     0,     0,     0,     0,
             0,     0,     0,     0,     0,     0],
        [31106, 33370,  5629, 16927,    54,    56, 31179, 11002,    72, 31526,
         31342,  8835, 31221, 31423,  7156,    55, 31396, 31222, 33370, 31604,
            72, 31415, 31396, 31222, 33370,  1231, 31221, 31195,  6851, 11002,
            75, 31106, 33370,  5629, 16927,    54,    56, 31179, 11002,    72,
          8835, 31221, 31423,    55, 31396, 31222, 33370, 31604,    72,  2926,
         31415, 31396, 31222, 33370,  1231, 31221, 11254, 11002, 31380,  1522,
         31482, 11002, 31380, 31640, 31187, 31222, 33370, 31135, 31396, 31380,
            73,     5,    54,    56, 34399,    55, 31191,    60,     5, 31415,
         31396, 31222, 33370,  1231, 31221, 31195,    60, 31179, 11002,    73,
          2122,    72,  6787, 31161, 31415, 31396, 31222, 33370,  1231, 31221,
         31195,    60, 31179, 11002,    73,     2]])

Tips：

这里 pad 的 0 和 mask 的 0 需要区分，这里 padding 为 0 是因为我们定义了 pad_token = unk_token，而 unk_token = <unk> 经过 tokenizer encode 后得到的是 [0]。

code=tokenizer.encode(tokenizer.pad_token, max_length=10000, truncation=True)
print('pad', tokenizer.unk_token, 'token', code)
=> pad <unk> token [0]

• seq_len

这个很好理解，一个 batch 内的多个样本，每个样本的 prompt_id 的长度即 Q_id 的长度。这个长度主要用于 batch 内判断 longest 最长的序列是多长，从而对短的样本进行 padding，保证进入深度模型网络的样本维度一致。

'seq_len': tensor([13, 31])

• labels

'labels': tensor([[31106, 15875,  77,  57,  53, 31159,  53,  60,    58, 31135,
          6787,  6775,    81, 31106,  3908, 14313, 31640, 31257, 31481, 31742,
            72,    57,    53, 31640, 31187,    53,    60,    58,  9323, 31178,
            52,    79,    54,    59,    56,     2,  -100,  -100,  -100,  -100,
          -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,
          -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,
          -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,
          -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,
          -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,
          -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,
          -100,  -100,  -100,  -100,  -100,  -100],
        [31106, 33370,  5629, 16927,    54,    56, 31179, 11002,    72, 31526,
         31342,  8835, 31221, 31423,  7156,    55, 31396, 31222, 33370, 31604,
            72, 31415, 31396, 31222, 33370,  1231, 31221, 31195,  6851, 11002,
            75, 31106, 33370,  5629, 16927,    54,    56, 31179, 11002,    72,
          8835, 31221, 31423,    55, 31396, 31222, 33370, 31604,    72,  2926,
         31415, 31396, 31222, 33370,  1231, 31221, 11254, 11002, 31380,  1522,
         31482, 11002, 31380, 31640, 31187, 31222, 33370, 31135, 31396, 31380,
            73,     5,    54,    56, 34399,    55, 31191,    60,     5, 31415,
         31396, 31222, 33370,  1231, 31221, 31195,    60, 31179, 11002,    73,
          2122,    72,  6787, 31161, 31415, 31396, 31222, 33370,  1231, 31221,
         31195,    60, 31179, 11002,    73,     2]])}

最后一个 key 是 labels，这里也可以解决我们前面的疑惑了，为什么样本里只有 input_ids 和 seq_len，但是经过 data_collator 处理送到 trainer 可以正常训练，这里的 label 对应的是 QA 里 A 的 token ids 即 target ids，以第一个 sample 为例，第一个 sample 的 A ids 为：

[31106, 3908, 14313, 31640, 31257, 31481, 31742, 72, 57, 53, 31640, 31187, 53, 60, 58, 9323, 31178, 52, 79, 54, 59, 56] + [2]

最后的 [2] 与前面 <unk> 类似，eos_token 后编码为 [2]：

code=tokenizer.encode(tokenizer.eos_token, max_length=10000, truncation=True)
print('eos', tokenizer.eos_token, 'token', code)
=> eos </s> token [2]

三.生成样本 By DIY

上面使用 Transformer API 实现了 {"input_ids": xxx, "seq_len": xxx} 形式的数据解析与样本生成，下面我们参考上面逻辑自定义一版 data_collator。首先整理一下思路：

input_ids = Q_ids + A_ids + [SEP]_ids

seq_len = Q_ids.length

labels = if (max_length) A_ids else A_ids + padding * n

根据 A_ids 的长度判断是否是 longest，然后决定是否补齐 padding 从而生成对应 Label。

1.样本准备

# p: [31106, 15875, 77, 57, 53, 31159, 53, 60, 58, 31135, 6787, 6775, 81]
# t: [31106, 3908, 14313, 31640, 31257, 31481, 31742, 72, 57, 53, 31640, 31187, 53, 60, 58, 9323, 31178, 52, 79, 54, 59, 56]
sample1 = [31106, 15875, 77, 57, 53, 31159, 53, 60, 58, 31135, 6787, 6775, 81, 31106, 3908, 14313, 31640, 31257, 31481,
           31742, 72, 57, 53, 31640, 31187, 53, 60, 58, 9323, 31178, 52, 79, 54, 59, 56, 2]

# p: [31106, 33370, 5629, 16927, 54, 56, 31179, 11002, 72, 31526, 31342, 8835, 31221, 31423, 7156, 55, 31396, 31222, 33370, 31604, 72, 31415, 31396, 31222, 33370, 1231, 31221, 31195, 6851, 11002, 75]
# t: [31106, 33370, 5629, 16927, 54, 56, 31179, 11002, 72, 8835, 31221, 31423, 55, 31396, 31222, 33370, 31604, 72, 2926, 31415, 31396, 31222, 33370, 1231, 31221, 11254, 11002, 31380, 1522, 31482, 11002, 31380, 31640, 31187, 31222, 33370, 31135, 31396, 31380, 73, 5, 54, 56, 34399, 55, 31191, 60, 5, 31415, 31396, 31222, 33370, 1231, 31221, 31195, 60, 31179, 11002, 73, 2122, 72, 6787, 31161, 31415, 31396, 31222, 33370, 1231, 31221, 31195, 60, 31179, 11002, 73]
sample2 = [31106, 33370, 5629, 16927, 54, 56, 31179, 11002, 72, 31526, 31342, 8835, 31221, 31423, 7156, 55, 31396,
           31222, 33370, 31604, 72, 31415, 31396, 31222, 33370, 1231, 31221, 31195, 6851, 11002, 75, 31106, 33370, 5629,
           16927, 54, 56, 31179, 11002, 72, 8835, 31221, 31423, 55, 31396, 31222, 33370, 31604, 72, 2926, 31415, 31396,
           31222, 33370, 1231, 31221, 11254, 11002, 31380, 1522, 31482, 11002, 31380, 31640, 31187, 31222, 33370, 31135,
           31396, 31380, 73, 5, 54, 56, 34399, 55, 31191, 60, 5, 31415, 31396, 31222, 33370, 1231, 31221, 31195, 60,
           31179, 11002, 73, 2122, 72, 6787, 31161, 31415, 31396, 31222, 33370, 1231, 31221, 31195, 60, 31179, 11002,
           73, 2]

json1 = {"input_ids": sample1, "seq_len": 13}
json2 = {"input_ids": sample2, "seq_len": 31}

features = [json1, json2]

样本继续使用上面构造的 json1 和 json2 并统一放到 features 的 list 中。

2.data_colloator 实现

def data_collator(features: list) -> dict:
    # 序列长度: [36, 106]
    len_ids = [len(feature["input_ids"]) for feature in features]
    # 取最长的序列长度: 106
    longest = max(len_ids)
    input_ids = []
    labels_list = []
    # 降序排列
    for ids_l, feature in sorted(zip(len_ids, features), key=lambda x: -x[0]):
        ids = feature["input_ids"]  # tokenIds
        seq_len = feature["seq_len"]  # seqLen
        # len(prompt) x [-100] + Target + [longest - len(prompt)] * [-100]
        labels = ([-100] * seq_len + ids[seq_len:] + [-100] * (longest - ids_l))
        ids = ids + [pad_token_id] * (longest - ids_l)

        _ids = torch.LongTensor(ids)
        labels_list.append(torch.LongTensor(labels))
        input_ids.append(_ids)
    # tensor([[], []])
    input_ids = torch.stack(input_ids)
    labels = torch.stack(labels_list)
    return {
        "input_ids": input_ids,
        "labels": labels,
    }

data_collator 的逻辑主要是这一句：

# len(prompt) x [-100] + Target + [longest - len(prompt)] * [-100]
labels = ([-100] * seq_len + ids[seq_len:] + [-100] * (longest - ids_l))

即将 Q mask 进行掩码，A 即 Target 保持不变，最后根据 longest 的长度决定是否 padding，这里手动指定了 pad_token_id = 0，实际代码中可以使用 tokenizer 自动指定：

tokenizer.pad_token_id

除此之外，这里和上面 API 生成还有一个区别是是否 mask Q_ids，API 把 Q_ids 全部去掉，上面方法保留了 Q_ids 的位置，但是使用 -100 进行了 mask，实际场景下二者没有区别，只不过 DIY 的 labels 如果 Q_ids 很长则会占用很多无关空间。

3.使用自定义 data_colloator

    trainer = ModifiedTrainer(
        model=model,
        train_dataset=dataset,
        args=training_args,
        callbacks=[TensorBoardCallback(writer)], # 回调将训练情况写入 tensorboard
        data_collator=data_collator,
    )

定义好 Trainer 后，将上面定一个 data_collator 传给 data_collator 参数即可，不过常规情况下我们使用 API 即可，如果自己对样本和 label 的构建有自定义需求，则可以采用后者 DIY 的形式。

四.总结

经过上面的分析，对于样本的处理和 label 的生成流程我们大致清晰了，下面解释下上面的样本如何应用在 LM 大语言模型中以及自己理解的这样构造的含义。

input 为 QA，output 为 A。Input 和 Output 分别在 Embedding lookup 获取输入向量，添加 position 向量后进入 Multi-Head Attention，首先经过 Q、K、V 的 Linear 映射转换，随后经过 Encoder 和 Decoder 的 Transformer 结构，最终通过一个 Linear + Softmax 输出预测概率 logits，输出概率 logits 与 A_ids 对应的 multi_hot 向量计算 batch loss 并回传，这里会忽略 [-100] 的掩码，更完整的 loss 计算逻辑大家可以参考 model.loss 的实现。

    # 根据输入计算 Loss
    def compute_loss(self, model, inputs, return_outputs=False):
        return model(
            input_ids=inputs["input_ids"],
            labels=inputs["labels"],
        ).loss

一般情况下，这里 QA 也分别代表 Prompt 和 Target，上面这种样本和 label 的构造方式意在学习当给定 Prompt 的前提提示下，模型能够预测得到 Target，即语言生成。如果想要学习其他不同的模式，大家也可以根据需求自定义修改上面的 data_collator。