如何用Hugging Face模型高效计算句子级困惑度并批量选最优句？

问题背景 #

手里有大量文档，每篇约含10个句子，需要为每篇文档找出困惑度最高（即微调后因果语言模型损失值最大）的句子，采用Hugging Face的distilgpt2模型。逐句处理结果正确但速度极慢（需处理约2500万句子），尝试批量处理时遇到两个核心问题：

• 分词器批量处理需padding，但计算损失时pad token会参与计算，导致单句损失受批中最长句长度影响，结果失真
• 批量输入模型仅返回标量损失（疑似批量均值），无法获取每个样本的单独损失

解决方案 #

1. 屏蔽pad token对损失的干扰 #

利用模型的attention_mask参数，在计算损失时自动忽略pad token的贡献。Hugging Face的因果语言模型支持传入该参数，模型会跳过mask值为0的token的损失计算，确保仅统计有效token的损失。

2. 获取单样本独立损失 #

默认返回的批量损失是均值，要得到每个样本的单独损失，需：

• 先计算每个token的交叉熵损失（禁用损失的自动归约）
• 对每个样本的有效token（非pad部分）的损失求均值，得到该样本的单句损失

优化后的代码 #

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
import spacy

# 初始化组件
tokenizer = AutoTokenizer.from_pretrained("distilgpt2")
tokenizer.pad_token = tokenizer.eos_token
model = AutoModelForCausalLM.from_pretrained("clm-gpu/checkpoint-138000")
model = model.to("cuda")  # 移至GPU加速计算
segmenter = spacy.load('en_core_web_sm')

def batch_compute_losses(sentences_batch):
    """批量计算输入句子列表中每个句子的损失"""
    # 批量分词，返回attention_mask用于屏蔽pad token
    encodings = tokenizer(
        sentences_batch,
        return_tensors="pt",
        padding=True,
        truncation=True
    ).to(model.device)
    
    input_ids = encodings.input_ids
    attention_mask = encodings.attention_mask

    # 前向传播获取logits，禁用梯度计算节省资源
    with torch.no_grad():
        outputs = model(input_ids=input_ids, attention_mask=attention_mask)
        logits = outputs.logits

    # 计算每个token的交叉熵损失（不自动归约）
    shift_logits = logits[..., :-1, :].contiguous()
    shift_labels = input_ids[..., 1:].contiguous()
    loss_fct = torch.nn.CrossEntropyLoss(reduction='none')
    token_losses = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
    
    # 还原为样本维度，仅统计有效token的损失均值
    token_losses = token_losses.view(shift_logits.size(0), shift_logits.size(1))
    valid_mask = attention_mask[..., 1:].contiguous()  # 对应预测token的有效mask
    sample_losses = (token_losses * valid_mask).sum(dim=1) / valid_mask.sum(dim=1)
    
    return sample_losses.cpu().numpy()

def process_documents(documents, batch_size=64):
    """批量处理所有文档，为每篇选出损失最大的句子"""
    all_best_sentences = []
    # 先整理所有句子及其所属文档索引
    sentence_to_doc = []
    all_sentences = []
    for doc_idx, document in enumerate(documents):
        sentences = [sentence.text.strip() for sentence in segmenter(document).sents]
        all_sentences.extend(sentences)
        sentence_to_doc.extend([doc_idx] * len(sentences))
    
    # 批量计算所有句子的损失
    num_sentences = len(all_sentences)
    losses = []
    for i in range(0, num_sentences, batch_size):
        batch_sentences = all_sentences[i:i+batch_size]
        batch_losses = batch_compute_losses(batch_sentences)
        losses.extend(batch_losses)
    
    # 按文档分组，筛选每篇的最优句子
    doc_sentences = {}
    doc_losses = {}
    for idx, (sentence, doc_idx, loss) in enumerate(zip(all_sentences, sentence_to_doc, losses)):
        if doc_idx not in doc_sentences:
            doc_sentences[doc_idx] = []
            doc_losses[doc_idx] = []
        doc_sentences[doc_idx].append(sentence)
        doc_losses[doc_idx].append(loss)
    
    for doc_idx in doc_sentences:
        max_loss_idx = doc_losses[doc_idx].index(max(doc_losses[doc_idx]))
        best_sentence = doc_sentences[doc_idx][max_loss_idx]
        all_best_sentences.append(best_sentence)
        # 这里可直接写入文件，或统一收集后批量写入
        # write(best_sentence)
    
    return all_best_sentences

# 使用示例
# documents = [...]  # 你的文档列表
# best_sentences = process_documents(documents)

核心优化点 #

• GPU加速：将模型和计算数据移至GPU，大幅提升运算效率
• 无梯度计算：通过torch.no_grad()禁用梯度追踪，节省显存与计算资源
• 大批次处理：支持跨文档的大批次计算，避免逐句处理的额外开销
• 精准损失计算：利用attention_mask屏蔽pad token，确保单句损失的准确性
• 高效分组筛选：先批量计算所有句子损失，再按文档分组选最优，减少重复操作

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