如何向Transformers Pipeline输入文本时限制为512个token？

问题场景 #

使用Transformers Pipeline测试本地预训练BERT模型时，短文本（token数<512）运行正常，但长文本触发报错：
RuntimeError: The size of tensor a (1593) must match the size of tensor b (512) at non-singleton dimension 1

现有代码 #

预测逻辑代码 #

def get_predicted_folder(text, model):
    pipe = pipeline("text-classification", model=model)
    if text:
        predicted_folder = pipe(text)
        label = predicted_folder[0]['label']
        score = predicted_folder[0]['score']
        return label, score
    else:
        err = "Error: The provided text is empty."
        return err, None

my_saved_model = "model/danish_bert_model"  # 本地保存的模型
label, score = get_predicted_folder(text, my_saved_model)

训练阶段分词代码 #

def tokenize_dataset(tokenizer, examples):
    return tokenizer(examples['text'], truncation=True, max_length=512)
tokenizer = AutoTokenizer.from_pretrained(model_ckpt)
dataset = dataset.map(lambda examples: tokenize_dataset(tokenizer, examples), batched=True)

已尝试的无效操作 #

• 为Pipeline传入tokenizer=model参数
• 提前加载分词器但未配置截断参数

解决方案 #

问题核心是Pipeline默认不对长文本做截断，而BERT模型最大输入长度固定为512，必须显式配置分词器的截断规则：

修正后的预测函数 #

from transformers import AutoTokenizer

def get_predicted_folder(text, model_path):
    # 加载与模型匹配的分词器
    tokenizer = AutoTokenizer.from_pretrained(model_path)
    # 初始化Pipeline时传入分词器并强制截断
    pipe = pipeline(
        "text-classification",
        model=model_path,
        tokenizer=tokenizer,
        truncation=True,
        max_length=512
    )
    if text:
        predicted_folder = pipe(text)
        return predicted_folder[0]['label'], predicted_folder[0]['score']
    else:
        return "Error: The provided text is empty.", None

my_saved_model = "model/danish_bert_model"
label, score = get_predicted_folder(text, my_saved_model)

关键要点 #

• 必须单独加载分词器实例传入Pipeline，不能直接传模型路径作为tokenizer参数
• 通过truncation=True和max_length=512确保输入文本被截断到模型支持的最大长度
• 若需保留更多文本信息，可实现滑动窗口分段预测后取加权平均，但基础场景下截断即可解决报错

内容的提问来源于stack exchange，提问作者D. F. R. 123