如何从HuggingFace的FeatureExtractionPipeline获取返回的attention mask？

在单节点8GPU环境下，使用HuggingFace的FeatureExtractionPipeline调用BAAI/bge-large-en-v1.5嵌入模型处理wikitext-2-v1数据集时，需获取返回的attention_mask来完成均值池化操作，核心实现代码如下：

from accelerate import Accelerator
from transformers import AutoTokenizer, AutoModel
import torch
from datasets import load_dataset
from transformers import pipeline

accelerator = Accelerator()

model_name = "BAAI/bge-large-en-v1.5"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModel.from_pretrained(model_name)

# 初始化特征抽取Pipeline，指定返回attention_mask
pipe = pipeline(
    "feature-extraction",
    model=model,
    tokenizer=tokenizer,
    max_length=512,
    truncation=True,
    padding=True,
    pad_to_max_length=True,
    batch_size=256,
    framework="pt",
    return_tensors=True,
    return_attention_mask=True,
    device=accelerator.device
)

# 加载数据集
dataset = load_dataset(
    "wikitext",
    "wikitext-2-v1",
    split="train",
)

# 均值池化函数：利用attention_mask忽略padding token的嵌入
def mean_pooling(model_output, attention_mask):
    token_embeddings = model_output[0]  # model_output的第一个元素是所有token的嵌入
    input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
    return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min=1e-9)

# 将数据集拆分到各GPU进程处理
with accelerator.split_between_processes(dataset["text"]) as data:
    for out in pipe(data):
        # 使用返回的attention_mask执行均值池化
        sentence_embeddings = mean_pooling(out, out["attention_mask"])

关键说明 #

• 通过Accelerator实现单节点多GPU的分布式推理调度
• 初始化Pipeline时必须设置return_attention_mask=True，才能在输出中获取注意力掩码
• accelerator.split_between_processes自动将数据集均分至各个GPU进程，避免数据重复处理
• 均值池化过程中，attention_mask会将padding token对应的嵌入值置0，确保池化结果仅为有效token的均值

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