如何适配transformers预训练分词器到TensorFlow翻译教程实现中英翻译

报错原因 #

你遇到的报错核心原因是：tf.data.Dataset管道在执行map操作时，传入tokenize_pairs的zh、en参数都是TensorFlow字符串张量类型，而HuggingFace Transformers库的分词器默认只接收Python原生字符串/字符串列表作为输入，张量类型没有strip方法，因此触发属性错误。

解决方案 #

你只需要用tf.py_function将原生Python分词逻辑包装为TensorFlow可识别的图操作，就能适配现有make_batches的流程，不需要修改原有批次生成逻辑，修改后的完整代码如下：

首先是分词和预处理部分的代码：

from transformers import BertTokenizer
import tensorflow as tf

# 加载分词器
tokenizer_en = BertTokenizer.from_pretrained("bert-base-cased")
tokenizer_zh = BertTokenizer.from_pretrained("bert-base-chinese")

# 定义原生Python分词逻辑，输入是张量转成的numpy数组
def preprocess_fn(zh_texts, en_texts):
    # 把字节数组转为字符串列表
    zh_texts = [b.decode('utf-8') for b in zh_texts.numpy()]
    en_texts = [b.decode('utf-8') for b in en_texts.numpy()]
    
    # 直接调用分词器接口，自动补全padding、返回tf张量
    zh_tokens = tokenizer_zh(zh_texts, padding='longest', return_tensors='tf', truncation=True, max_length=128)['input_ids']
    en_tokens = tokenizer_en(en_texts, padding='longest', return_tensors='tf', truncation=True, max_length=128)['input_ids']
    
    return zh_tokens, en_tokens

# 包装为tf.data管道可调用的函数
def tokenize_pairs(zh, en):
    zh_tokens, en_tokens = tf.py_function(
        func=preprocess_fn,
        inp=[zh, en],
        Tout=[tf.int32, tf.int32]
    )
    # 固定张量形状，避免后续图执行阶段报错
    zh_tokens.set_shape([None, 128])
    en_tokens.set_shape([None, 128])
    return zh_tokens, en_tokens

原有make_batches函数不需要做任何修改，直接运行即可正常生成训练和验证批次。

注意事项 #

• 代码里的max_length可根据你的任务需求调整，只要中英序列的长度设置匹配即可
• 如果不需要截断可以去掉truncation=True参数，但要注意数据集中最长文本长度不要超出显存承载范围

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