使用Hugging Face Transformers时如何显示文本摘要进度百分比？

一、添加进度显示 #

transformers的summarization pipeline默认无进度提示，可拆分tokenizer与model手动实现，结合tqdm跟踪生成进度：

import torch
from tqdm import tqdm
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

wall_of_text = "some long text..."

# 加载tokenizer与模型，启用FP16加速
tokenizer = AutoTokenizer.from_pretrained("pszemraj/long-t5-tglobal-base-16384-book-summary")
model = AutoModelForSeq2SeqLM.from_pretrained(
    "pszemraj/long-t5-tglobal-base-16384-book-summary",
    torch_dtype=torch.float16 if torch.cuda.is_available() else None,
    device_map="auto"
)

# 编码输入文本
inputs = tokenizer(
    wall_of_text,
    return_tensors="pt",
    max_length=16384,
    truncation=True
).to(model.device)

# 设置生成参数
generate_kwargs = {
    "min_length": 500,
    "max_length": 2000,  # 按需调整，无需设到模型上限16384
    "no_repeat_ngram_size": 3,
    "encoder_no_repeat_ngram_size": 3,
    "repetition_penalty": 3.5,
    "num_beams": 2,  # 减少beam数量提升速度
    "early_stopping": True,
}

# 初始化进度条
total_tokens = generate_kwargs["max_length"]
progress_bar = tqdm(total=total_tokens, desc="生成摘要")

# 自定义回调更新进度
def update_progress(step, outputs):
    progress_bar.update(1)

# 执行生成
outputs = model.generate(
    **inputs,
    **generate_kwargs,
    callback=update_progress
)

# 解码并保存结果
summary_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
with open('pszemraj-long-t5-tglobal-base-16384-book-summary.txt', 'w') as f:
    f.write(summary_text)

progress_bar.close()

二、优化RTX 3060运行速度 #

针对9k token输入耗时过长问题，可通过以下方式优化：

• 启用FP16精度：加载模型时指定torch_dtype=torch.float16，RTX 3060支持FP16运算，能大幅降低显存占用并提升速度。
• 减少beam数量：num_beams=4会带来4倍计算量，降至num_beams=2或使用贪心搜索（num_beams=1），速度提升明显，效果差异可接受。
• 降低max_length：无需将摘要长度设到模型上限16384，根据输入文本长度调整到合理范围（如输入9k token时设为1000-2000），减少生成token数直接缩短耗时。
• 分块处理长文本：若输入文本远超16k token，可拆分为多个16k以内的块，分别生成摘要后再合并，避免单块处理的显存压力与长时间等待。

内容的提问来源于stack exchange，提问作者Furkan Gözükara