如何向Hugging Face Transformers的pipeline传递参数获取固定尺寸BERT embedding

核心原因 #

你之前的方案不生效，是因为Transformers的pipeline内部调用tokenizer时，会默认覆盖你在tokenizer初始化时设置的padding、截断相关参数，不会直接复用初始化时的配置，所以需要主动将预处理参数传递给pipeline的预处理流程。

正确实现方式 #

方案1：每次调用pipeline时传递参数 #

适合需要动态调整max_length等参数的场景，调用时直接将tokenizer的参数传入即可：

from transformers import AutoTokenizer, AutoModel, pipeline

# 基础初始化tokenizer和模型
tokenizer = AutoTokenizer.from_pretrained("emilyalsentzer/Bio_ClinicalBERT")
model = AutoModel.from_pretrained("emilyalsentzer/Bio_ClinicalBERT")

# 初始化pipeline
nlp = pipeline("feature-extraction", model=model, tokenizer=tokenizer, device=0)

# 调用时传入预处理参数
embedding = nlp(
    "hello world",
    padding="max_length",
    truncation=True,
    max_length=500
)
# 此时输出的embedding长度固定为500

方案2：初始化pipeline时设置全局预处理参数 #

适合固定参数的场景，初始化时通过preprocess_kwargs指定默认预处理参数，后续调用无需重复传参：

nlp = pipeline(
    "feature-extraction",
    model=model,
    tokenizer=tokenizer,
    device=0,
    # 指定预处理阶段的默认参数
    preprocess_kwargs={
        "padding": "max_length",
        "truncation": True,
        "max_length": 500
    }
)

# 直接调用即可得到固定尺寸输出
embedding = nlp("hello world")

输出验证 #

你可以通过以下代码确认输出尺寸符合预期：

import numpy as np
print(np.array(embedding).shape)
# 预期输出：(1, 500, 768) （768为Bio_ClinicalBERT的隐藏层维度）

内容的提问来源于stack exchange，提问作者Israel-abebe