使用Transformers Tokenizer批量编码对话数据集报错的解决方法

问题背景 #

你的数据集为嵌套列表格式，外层元素代表一段对话，内层元素是对话中的单轮utterance：

foo = [['asdf', 'hello', 'john matthew'],
       ['asdf', 'good bye', 'luke caleb'],
       ['asdf', 'see you tomorrow', 'peter simon']]

使用Transformers Tokenizer的batch_encode_plus方法编码时触发错误：

TypeError: TextEncodeInput must be Union[TextInputSequence, Tuple[InputSequence, InputSequence]]

文档标注该方法支持typing.List[typing.List[str]]类型，但实际报错，需解决此问题并找到适配PyTorch Lightning的高效编码方式。

错误原因 #

batch_encode_plus支持的List[List[str]]是文本对格式（每个子列表是一组配对文本，如[查询句, 文档句]），而非多轮对话的嵌套列表。你的数据是多轮utterance组成的对话，不属于文本对范畴，因此类型不匹配。

解决方法 #

1. 拼接多轮对话为单段文本编码 #

如果模型将整段对话作为单一输入，先把每段对话的所有utterance拼接成单个字符串，再传入编码方法：

from transformers import AutoTokenizer
import torch

tokenizer = AutoTokenizer.from_pretrained('bert-base-uncased', use_fast=True)

# 用tokenizer自带的分隔符拼接每轮utterance，可根据模型需求调整
processed_foo = [tokenizer.sep_token.join(turns) for turns in foo]

bar = tokenizer.batch_encode_plus(
    processed_foo,
    return_tensors='pt',
    padding='max_length',
    truncation=True,
    max_length=50
)

# 输出形状：torch.Size([3, 50])，对应3条对话，每条编码后长度50
print(bar['input_ids'].shape)

2. 保留多轮结构的独立编码 #

若需要保留每轮utterance的独立编码（用于对话历史建模），可遍历对话对每轮单独编码，再统一对齐长度：

encoded_dialogues = []
single_turn_max_len = 15  # 单轮utterance的最大长度，按需调整

for dialogue in foo:
    # 编码单轮对话
    encoded_turns = tokenizer(
        dialogue,
        padding='max_length',
        truncation=True,
        max_length=single_turn_max_len,
        return_tensors='pt'
    )
    # 堆叠每轮编码结果，形成[轮数, 单轮长度]的张量
    dialogue_input_ids = torch.stack(encoded_turns['input_ids'], dim=0)
    dialogue_attention_mask = torch.stack(encoded_turns['attention_mask'], dim=0)
    encoded_dialogues.append({
        'input_ids': dialogue_input_ids,
        'attention_mask': dialogue_attention_mask
    })

# 对齐所有对话的轮数（补零到最多轮数）
max_turns = max(len(dialogue) for dialogue in foo)
for item in encoded_dialogues:
    item['input_ids'] = torch.nn.functional.pad(
        item['input_ids'], (0, 0, 0, max_turns - item['input_ids'].shape[0])
    )
    item['attention_mask'] = torch.nn.functional.pad(
        item['attention_mask'], (0, 0, 0, max_turns - item['attention_mask'].shape[0])
    )

# 合并为batch张量，形状为[batch_size, 轮数, 单轮长度]
batch_input_ids = torch.stack([d['input_ids'] for d in encoded_dialogues], dim=0)
batch_attention_mask = torch.stack([d['attention_mask'] for d in encoded_dialogues], dim=0)
print(batch_input_ids.shape)  # torch.Size([3, 3, 15])

3. 适配PyTorch Lightning的高效批量处理 #

推荐使用Hugging Face Dataset类配合map方法，实现多进程批量编码，直接适配PyTorch Lightning数据模块：

from datasets import Dataset

# 转换为Dataset格式
dataset = Dataset.from_dict({'dialogue': foo})

# 定义编码函数
def encode_dialogues(examples):
    # 拼接对话为单段文本（需保留多轮结构可参考方法2修改逻辑）
    texts = [tokenizer.sep_token.join(turns) for turns in examples['dialogue']]
    return tokenizer(
        texts,
        padding='max_length',
        truncation=True,
        max_length=50,
        return_tensors='pt'
    )

# 多进程批量编码，num_proc指定进程数加速
encoded_dataset = dataset.map(encode_dialogues, batched=True, num_proc=4)

# 设置为PyTorch张量格式，可直接传入DataLoader
encoded_dataset.set_format(type='torch', columns=['input_ids', 'attention_mask'])

内容的提问来源于stack exchange，提问作者loose the fools juice