别再只盯着input_ids了！用PyTorch和Transformers处理文本时，attention_mask和token_type

别再只盯着input_ids了！用PyTorch和Transformers处理文本时，attention_mask和token_type_ids的实战避坑指南

当你第一次接触Transformer模型时，input_ids可能是最吸引你注意力的部分——毕竟它直观地展示了文本如何被转换为数字。但很快你会发现，仅仅理解input_ids是远远不够的。在实际项目中，特别是在处理不等长序列、句子对任务或需要精确控制模型注意力的场景中，attention_mask和token_type_ids这两个看似"配角"的字段，往往成为决定模型表现的关键因素。

1. 为什么你需要关注attention_mask和token_type_ids

在构建问答系统或文本分类模型时，我们经常遇到这样的困惑：明明输入的数据看起来没问题，模型训练却出现异常波动，或者推理时产生不符合预期的结果。这些问题很多时候都源于对attention_mask和token_type_ids的误解或不当使用。

attention_mask本质上是一个二进制张量，它告诉模型哪些token是真实的文本内容，哪些是填充的无效token。想象一下，当你把不同长度的句子批量输入模型时，较短的句子会被填充(padding)到与最长句子相同的长度。如果没有attention_mask，模型会平等对待所有token——包括那些无意义的填充token，这显然会干扰模型的学习。

token_type_ids则主要用于区分句子对中的不同部分。在问答任务中，它帮助模型识别哪些token属于问题，哪些属于上下文；在文本对分类任务中，它标记两个句子的边界。忽视这个字段可能导致模型混淆句子关系，严重影响性能。

实际案例：在一个客户支持的工单分类项目中，团队发现模型对长文本的分类准确率显著低于短文本。排查后发现是因为没有正确设置attention_mask，导致模型被大量填充token干扰。

2. attention_mask的深度解析与实战技巧

2.1 attention_mask的工作原理

attention_mask是一个与input_ids形状相同的张量，其中：

1表示对应的token是有效内容
0表示对应的token是填充内容

在Transformer的自注意力机制中，attention_mask直接影响注意力权重的计算。具体来说，对于被标记为0的位置，模型会赋予极小的注意力分数（如-10000），使得这些位置几乎不影响最终的表示。

from transformers import AutoTokenizer tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased") # 两个长度不同的句子 sentences = ["Hello, world!", "This is a longer sentence for demonstration."] # 自动填充并生成attention_mask inputs = tokenizer(sentences, padding=True, return_tensors="pt") print("Input IDs:\n", inputs["input_ids"]) print("Attention Mask:\n", inputs["attention_mask"])

2.2 常见陷阱与解决方案

陷阱1：手动填充导致attention_mask不匹配

开发者有时会先对文本进行手动填充，再调用tokenizer。这种做法会导致attention_mask无法正确反映实际的填充位置。

# 不推荐的做法 padded_text = ["Hello, world! [PAD] [PAD]", "This is a longer sentence."] inputs = tokenizer(padded_text, return_tensors="pt") # attention_mask将错误标记 # 正确的做法 raw_text = ["Hello, world!", "This is a longer sentence."] inputs = tokenizer(raw_text, padding=True, return_tensors="pt")

陷阱2：忽略attention_mask在自定义模型中的传递

当你基于Transformer架构构建自定义模型时，必须确保attention_mask被正确传递到各层：

import torch from transformers import AutoModel model = AutoModel.from_pretrained("bert-base-uncased") outputs = model( input_ids=inputs["input_ids"], attention_mask=inputs["attention_mask"] # 必须显式传递 )

2.3 高级应用：动态注意力控制

attention_mask不仅可以用于标记填充，还能实现更精细的注意力控制。例如，在生成任务中，你可以防止模型"偷看"未来的token：

# 创建因果注意力mask (用于自回归生成) seq_length = inputs["input_ids"].shape[1] causal_mask = torch.tril(torch.ones(seq_length, seq_length)).bool() # 结合padding mask和causal mask combined_mask = inputs["attention_mask"].unsqueeze(1) & causal_mask

3. token_type_ids的实战应用

3.1 何时需要token_type_ids

token_type_ids主要用于以下场景：

问答系统（问题 vs 上下文）
句子对分类（如自然语言推理）
任何需要模型区分不同文本片段的任务

# 句子对示例 question = "What is the capital of France?" context = "Paris is the capital of France." inputs = tokenizer(question, context, return_tensors="pt") print("Token type IDs:\n", inputs["token_type_ids"])

3.2 常见错误排查

错误1：错误拼接导致句子边界混淆

手动拼接句子而不使用tokenizer的句子对接口，会导致token_type_ids失效：

# 错误做法 wrong_input = tokenizer(question + " " + context) # 丢失句子边界信息 # 正确做法 correct_input = tokenizer(question, context) # 自动生成token_type_ids

错误2：忽略模型是否实际使用token_type_ids

并非所有Transformer模型都使用token_type_ids。例如，RoBERTa就不依赖这个字段。使用前应检查模型文档：

from transformers import RobertaTokenizer roberta_tokenizer = RobertaTokenizer.from_pretrained("roberta-base") inputs = roberta_tokenizer(question, context) # 不会返回token_type_ids

3.3 多片段场景的扩展应用

对于需要三个或更多文本片段的任务（如多文档问答），可以通过扩展token_type_ids来实现：

doc1 = "Paris is the capital of France." doc2 = "It has many famous landmarks." question = "What is the capital of France?" # 自定义处理 encoding1 = tokenizer(doc1, return_tensors="pt") encoding2 = tokenizer(doc2, add_special_tokens=False, return_tensors="pt") encoding_q = tokenizer(question, add_special_tokens=False, return_tensors="pt") # 手动拼接和创建token_type_ids input_ids = torch.cat([ encoding1["input_ids"], encoding2["input_ids"], torch.tensor([tokenizer.sep_token_id]), encoding_q["input_ids"] ], dim=1) token_type_ids = torch.cat([ torch.zeros_like(encoding1["input_ids"]), torch.ones_like(encoding2["input_ids"]), torch.ones_like(torch.tensor([[tokenizer.sep_token_id]])), torch.zeros_like(encoding_q["input_ids"]) + 2 # 使用2表示问题 ], dim=1)

4. 综合案例：构建一个健壮的问答系统管道

让我们将这些知识应用到一个完整的问答系统实现中。这个案例将展示如何正确处理attention_mask和token_type_ids，以及如何避免常见的性能陷阱。

4.1 数据预处理管道

from transformers import AutoTokenizer, AutoModelForQuestionAnswering import torch tokenizer = AutoTokenizer.from_pretrained("bert-large-uncased-whole-word-masking-finetuned-squad") model = AutoModelForQuestionAnswering.from_pretrained("bert-large-uncased-whole-word-masking-finetuned-squad") def preprocess_qa_examples(questions, contexts, max_length=384, stride=128): """ 处理问答对，支持长文档的分块处理 """ inputs = tokenizer( questions, contexts, max_length=max_length, truncation="only_second", stride=stride, return_overflowing_tokens=True, return_offsets_mapping=True, padding="max_length", return_tensors="pt" ) # 对于分块处理，需要调整offset_mapping sample_mapping = inputs.pop("overflow_to_sample_mapping") offset_mapping = inputs.pop("offset_mapping") # 标记无法回答的问题 inputs["start_positions"] = [] inputs["end_positions"] = [] for i, offset in enumerate(offset_mapping): # 这里可以添加真实答案位置的标记逻辑 inputs["start_positions"].append(0) inputs["end_positions"].append(0) return inputs

4.2 推理过程中的注意力优化

在推理时，合理利用attention_mask可以显著提升效率：

def predict_answer(question, context, model, tokenizer): inputs = tokenizer( question, context, max_length=512, truncation="only_second", return_tensors="pt", padding="max_length" ) # 创建更精细的attention_mask # 减少对[SEP]之后padding的计算 sep_index = (inputs["input_ids"0] == tokenizer.sep_token_id).nonzero()[0, 0].item() inputs["attention_mask"0, sep_index+1:] = 0 with torch.no_grad(): outputs = model(**inputs) answer_start = torch.argmax(outputs.start_logits) answer_end = torch.argmax(outputs.end_logits) + 1 answer = tokenizer.convert_tokens_to_string( tokenizer.convert_ids_to_tokens( inputs["input_ids"0][answer_start:answer_end] ) ) return answer