源代码来自：https://github.com/huggingface/nanoVLM/blob/main/data/collators.py
详解如下所示：
import torch#-------------------------------#
# 主要是在数据加载器的构建中被使用
#-------------------------------#class BaseCollator(object):def __init__(self, tokenizer):self.tokenizer = tokenizerrandom_string_5_letters = "xzyvd" # 作为“锚点”，查找它在模板化后的完整文本中的位置# 将输入消息转换成Chat模板格式的字符串 例如 "<|start|>assistant\nxzyvd<|end|>" 此为纯文本而不是被编码后得到的token idsrandom_string_chat_templated = self.tokenizer.apply_chat_template([{"role": "assistant", "content": random_string_5_letters}], tokenize=False, add_special_tokens=False)random_string_location = random_string_chat_templated.find(random_string_5_letters) # 查找我们之前插入的“随机标记”出现的位置# 例如回复为<|start|>assistant\nxzyvd<|end|># 获取到nxzyvd开始后的位置, 然后从而获取到前缀的长度# 目的是在后续设置loss_mask时能够精准跳过模板前缀，只对assistant回复的实际内容进行监督self.prefix_len = len(self.tokenizer.encode(random_string_chat_templated[:random_string_location])) # 找到前缀模板结束的位置#----------------------------------------------------------## 用于处理批量对话消息# 随后返回模型需要的token ids、attention mask以及loss mask# 1.将消息转换为模型所需的 token 格式# 2.根据消息中的role（例如 assistant）标记哪些token需要计算损失（loss_mask），即只对assistant的具体输出进行损失计算，而不对user的内容进行计算# 3.将所有输入统一padding到最大长度max_len，确保批次的输入大小一致#----------------------------------------------------------#def prepare_inputs_and_loss_mask(self, batched_messages, max_length=None):batch_token_ids: list[list[int]] = [] # 保存每个批次消息的token idsbatch_masks: list[list[int]] = [] # 保存每个批次消息的loss_mask，即哪些token需要计算损失batch_attentions: list[list[int]] = [] # 保存每个批次消息的attention mask，模型用来指示哪些部分是有效输入，哪些是 paddingfor messages in batched_messages: # 每一条消息中都包含若干user和assistant的内容#---------------------------------------------------------------------------------------## 对于此处生成的attention mask# tokenizer会自动将padding部分的attention mask设为0，其余为1# 其作用为告诉模型哪些token是“真正需要注意的内容”，哪些只是为了凑长度而padding的垃圾位# 它是Transformer中注意力机制不可或缺的一部分，尤其在处理变长输入（如自然语言对话）时非常关键# NOTE：此处，tokenizer没有做统一长度 padding，而是保留了变长的attention_mask#---------------------------------------------------------------------------------------#conv_ids = self.tokenizer.apply_chat_template(messages,tokenize=True, # 控制attention mask相关内容add_special_tokens=False,return_dict=True,) # conv_ids是面向整个对话的一个字典，包含了对应的 input_ids（token ids）和 attention_maskmask = [0] * len(conv_ids["input_ids"]) # 为每个对话消息初始化一个全零的 mask 列表# Locate each assistant turn and flip its mask to 1cursor = 0 # 用来记录当前已经处理过的token数量for msg in messages: # 对user与assistant的内容均进行处理segment_ids = self.tokenizer.apply_chat_template([msg], tokenize=True, add_special_tokens=False) # 将每条消息msg转换为token ids # 只包含这一条消息的内容seg_len = len(segment_ids) # 获取消息的长度, 即为每条消息的实际token数目#---------------------------------------## 当处理角色为assistant的时候展开下述操作# 只对其具体回复的内容进行操作#---------------------------------------#if msg["role"] == "assistant":start = cursor + self.prefix_len # 确定消息的起点end = cursor + seg_len # 根据消息的长度去确定终点mask[start:end] = [1] * (end - start) # attend to these tokens # 将assistant的回复部分的mask设置为1cursor += seg_len # 因为一组对话中assistant回复的内容可能有多处, 因此需要进行累积batch_token_ids.append(conv_ids["input_ids"]) # token idsbatch_masks.append(mask) # 哪些token需要去计算batch_attentions.append(conv_ids["attention_mask"]) # 哪些部分是有效输入# NOTE：主要针对assistant回复过长的情况进行处理if max_length is not None: # We need to keep the tokens to allow for the img embed replacing logic to work. Otherwise, we would need to track which images correspond to long samples.batch_token_ids = [ids[:max_length] for ids in batch_token_ids] # 对超过max length的样本进行裁剪, 使其长度满足要求# 如果长度超过 max_length，则将其截断为全零的 mask（表示忽略该样本）batch_masks = [m if len(m) <= max_length else [0]*max_length for m in batch_masks] # Ignore samples that are longer than max_lengthbatch_attentions = [a[:max_length] for a in batch_attentions] # 同样进行截取# Pad samples to max lengthif max_length is not None:max_len = max_lengthelse:max_len = max(map(len, batch_token_ids))# 对每个样本均展开padding操作batch_token_ids = [[self.tokenizer.pad_token_id]*(max_len-len(ids)) + ids for ids in batch_token_ids] # 使用pad_token_id将长度填充到max lengthbatch_masks = [[0]*(max_len-len(m)) + m for m in batch_masks] # 填充至最大长度max_len，使用0填充batch_attentions = [[0]*(max_len-len(a)) + a for a in batch_attentions] # 填充至最大长度max_len，使用0填充 # NOTE: 相当于是在tokenzier的基础上 根据max length去展开补充性paddingreturn torch.tensor(batch_token_ids), torch.tensor(batch_attentions), torch.tensor(batch_masks).to(torch.bool)#-------------------------------------#
# Visual Question Answering Collator
# 训练与验证数据集
#-------------------------------------#
class VQACollator(BaseCollator):def __init__(self, tokenizer, max_length):self.max_length = max_lengthsuper().__init__(tokenizer)def __call__(self, batch):images = [item["images"] for item in batch]messages_batched = [item["text_data"] for item in batch]# Stack imagesimgs = [img for sublist in images for img in sublist]images = torch.stack(imgs)# Create inputs by concatenating special image tokens, question, and answerbatch_input_ids, batch_attention_mask, loss_masks = self.prepare_inputs_and_loss_mask(messages_batched, max_length=self.max_length)#--------------------------------------------------------------------------------------------------------------------------------------------------------------------------## Create labels where only answer tokens are predicted# 1. 首先将模型回复的内容全部复制一份出来, 然后将为mask为0的区域全部填充为-100, 表明直接忽视不参与计算# 2. 为适应因果语言建模, 展开标签平移操作, 作用为确保模型在展开语言生成任务时, 能够预测当前时间步的下一个token# 具体而言, labels[:, :-1]为选择每个样本的所有token中除去最后一个token的部分, labels[:, 1:]为获取每个样本中从第二个token到最后一个token的所有内容# 这样就可以将每个样本的所有token都可以向左移动一位, 从而将每个位置对应的token都用它的下一个token去进行预测。这样每个token的标签都变成了它的下一个token, 即为next token prediction# 3. 这样最后一个token由于没有标签目标, 直接设置为-100即可, 表明到了结尾# 例子：# batch_input_ids为[[101, 2001, 2023, 2045, 102]], 其中2001处的loss mask为0, 那么labels即为[[101, 2023, 2045, 102]]# 然后第一个样本的0 1 2 3四个位置上对应的label即变为[2023, 2045, 102, -100]# 这样就形成了真值标签[[2023, 2045, 102, -100]]#--------------------------------------------------------------------------------------------------------------------------------------------------------------------------#labels = batch_input_ids.clone().masked_fill(~loss_masks, -100) # 将~loss_masks为1的地方填充为-100 NOTE：此处相当于就是无效的地方labels[:, :-1] = labels[:, 1:] # Shift labels for causal LMlabels[:, -1] = -100 # Last token has no targetreturn {"image": images, # 图像"input_ids": batch_input_ids, # 输入内容"attention_mask": batch_attention_mask, # 告诉模型在等长序列中, 哪些是需要关注的实际token, 哪些是padding token"labels": labels, #标签}#--------------------------------------------------------#
# 测试数据集
# https://huggingface.co/datasets/Lin-Chen/MMStar
#--------------------------------------------------------#
class MMStarCollator(BaseCollator): def __init__(self, tokenizer):super().__init__(tokenizer)def __call__(self, batch):images = [item["image"] for item in batch]messages_batched = [item["text_data"] for item in batch]# Stack imagesimages = torch.stack(images)# Create inputs by concatenating special image tokens, question, and answerbatch_input_ids, batch_attention_mask, loss_masks = self.prepare_inputs_and_loss_mask(messages_batched)#---------------------------------------------------------------------------------------------------------------------------------------------## 1. 把需要预测的位置（即 loss_masks=1）设成pad token, 这意味着这些位置不会被送去模型作为“输入”，因为它们是模型需要生成的内容# 2. 把要预测的部分在attention mask里屏蔽掉, 导致模型不会“看到”这些 token，符合推理阶段的auto-regressive decoding 逻辑# 3. 只保留需要预测的token作为标签，其余地方用pad填充#---------------------------------------------------------------------------------------------------------------------------------------------#"""example:query: "User: What color is the sky?\nAssistant: The sky is"prediction: "blue."那么 loss_mask 会标记 "blue." 这一段, collator就会:把 input_ids 中 "blue." 变成pad(输入时忽略)把 attention_mask 中对应位置设为0(不关注)把 labels 中 "blue." 保留, 其余是pad(只评估蓝天这个词)"""input_ids = batch_input_ids.masked_fill(loss_masks, self.tokenizer.pad_token_id)attention_mask = batch_attention_mask.masked_fill(loss_masks, 0)labels = batch_input_ids.clone().masked_fill(~loss_masks, self.tokenizer.pad_token_id)return {"images": images,"input_ids": input_ids,"attention_mask": attention_mask,"labels": labels,}
