Token标注任务中：如何合并多份BatchEncoding或创建TensorFlow数据集？

在Token标注任务中，使用Transformers分词器时，由于需要处理子词拆分后的标签重排，只能逐文本处理得到BatchEncoding对象列表。以下是将该列表转换为TensorFlow Dataset或合并为单个BatchEncoding的两种可行方案：

可复现示例代码 #

先基于你的场景生成BatchEncoding列表：

from transformers import BertTokenizerFast
import tensorflow as tf

tokenizer = BertTokenizerFast.from_pretrained('bert-base-multilingual-uncased')
tokens = []
for text in ["Hello there", "Good morning"]:
    t = tokenizer(text.split(), is_split_into_words=True, padding='max_length', truncation=True, max_length=10)
    t['labels'] = list(map(lambda x: 1, t.word_ids()))  # 模拟标签生成逻辑
    tokens.append(t)

方案1：合并BatchEncoding列表为单个BatchEncoding #

每个BatchEncoding本质是包含input_ids、attention_mask、labels等键的字典，我们可以将相同键的所有值合并为二维数组，再创建新的BatchEncoding：

from transformers import BatchEncoding
import numpy as np

# 初始化合并后的字典
merged_data = {}
# 遍历每个键，收集所有样本对应的值并转为numpy数组
for key in tokens[0].keys():
    merged_data[key] = np.array([sample[key] for sample in tokens])

# 创建合并后的BatchEncoding对象
merged_batch = BatchEncoding(merged_data)

# 转换为TensorFlow Dataset
tf_dataset = tf.data.Dataset.from_tensor_slices(merged_batch)

方案2：直接将BatchEncoding列表转为TensorFlow Dataset #

无需合并，直接将每个BatchEncoding转为普通字典后，用from_tensor_slices创建数据集：

# 将每个BatchEncoding转为普通字典，组成列表
dataset_dicts = [dict(batch) for batch in tokens]

# 生成TensorFlow Dataset
tf_dataset = tf.data.Dataset.from_tensor_slices(dataset_dicts)

两种方案都能生成符合训练要求的TensorFlow Dataset，可直接接入模型训练流程。

内容的提问来源于stack exchange，提问作者raquelhortab