如何在Hugging Face预训练文本摘要模型中使用自定义Tokenizer避免重复分词

核心原理 #

Transformers的官方pipeline支持跳过内部分词步骤，只要满足两个前提：

• 提前分词时返回PyTorch/TensorFlow格式的张量结果
• 调用pipeline时传入tokenizer=None参数，同时自行完成长文本的token分片、补全操作

完整实现代码 #

1. 初始化与提前分词 #

你使用的t5-small模型做摘要任务需要先给原文加前缀summarize: ，如果使用其他模型可省略该步骤：

import transformers
import torch

# 初始化摘要pipeline
summarizer = transformers.pipeline("summarization", model = 't5-small', tokenizer = 't5-small')
# 拼接t5模型要求的任务前缀 + 待处理长文本
text = "summarize: " + "待处理的超长原始文本内容..."

# 提前完成全量分词，不做截断，后续自行分片
tokenized_text = summarizer.tokenizer(
    text,
    return_tensors="pt",
    truncation=False
)
# 去除多余的batch维度
input_ids = tokenized_text["input_ids"][0]
attention_mask = tokenized_text["attention_mask"][0]

2. 长token分片处理 #

# 读取模型支持的最大输入长度，t5-small为512，可根据实际使用的模型自动读取
max_input_len = summarizer.model.config.max_position_embeddings
# 分片重叠长度，避免分片处上下文断裂，可自行调整
stride = 50

chunks = []
# 遍历生成分片
for i in range(0, len(input_ids), max_input_len - stride):
    # 截取当前分片的token与注意力掩码
    chunk_input_ids = input_ids[i:i+max_input_len]
    chunk_attention_mask = attention_mask[i:i+max_input_len]
    # 对长度不足的最后一个分片做padding补齐
    padding_len = max_input_len - len(chunk_input_ids)
    if padding_len > 0:
        chunk_input_ids = torch.cat([chunk_input_ids, torch.tensor([summarizer.tokenizer.pad_token_id]*padding_len)])
        chunk_attention_mask = torch.cat([chunk_attention_mask, torch.tensor([0]*padding_len)])
    # 恢复batch维度，符合模型输入要求
    chunks.append({
        "input_ids": chunk_input_ids.unsqueeze(0),
        "attention_mask": chunk_attention_mask.unsqueeze(0)
    })

3. 跳过分词直接推理 #

chunk_summaries = []
for chunk in chunks:
    # 传入tokenizer=None跳过pipeline内部的分词步骤
    summary = summarizer(
        chunk,
        tokenizer=None,
        max_length=150, # 生成摘要的最大长度，可自定义
        min_length=30, # 生成摘要的最小长度，可自定义
        do_sample=False
    )[0]["summary_text"]
    chunk_summaries.append(summary)

# 合并所有分片的摘要结果，得到最终全文摘要
final_summary = " ".join(chunk_summaries)

注意事项 #

• 若使用GPU推理，可提前将token张量移动到CUDA设备，大幅提升处理速度
• 分片重叠长度stride可根据文本类型调整，专业类文本建议设置更大的重叠值，减少信息丢失
• 非T5类摘要模型不需要加summarize: 前缀，按对应模型的官方输入要求调整即可

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