如何正确使用HuggingFace生物医学模型获取可信文本嵌入？

问题描述 #

我需要为生物医学文本生成文本嵌入，尝试用sentence-transformers框架加载microsoft/BiomedNLP-PubMedBERT-base-uncased-abstract-fulltext模型，代码如下：

from sentence_transformers import SentenceTransformer
import pandas as pd

sbert_model = SentenceTransformer('microsoft/BiomedNLP-PubMedBERT-base-uncased-abstract-fulltext')
document_embeddings = sbert_model.encode(pd.Series(['hello', 'cell type', 'protein']))
print(document_embeddings)

运行后出现以下警告：

No sentence-transformers model found with name /home/user/.cache/torch/sentence_transformers/microsoft_BiomedNLP-PubMedBERT-base-uncased-abstract-fulltext. Creating a new one with MEAN pooling.
Some weights of the model checkpoint at /home/user/.cache/torch/sentence_transformers/microsoft_BiomedNLP-PubMedBERT-base-uncased-abstract-fulltext were not used when initializing BertModel: ['cls.predictions.transform.dense.weight', 'cls.predictions.transform.LayerNorm.bias', 'cls.seq_relationship.bias', 'cls.predictions.decoder.bias', 'cls.predictions.transform.dense.bias', 'cls.seq_relationship.weight', 'cls.predictions.decoder.weight', 'cls.predictions.transform.LayerNorm.weight', 'cls.predictions.bias']
- This IS expected if you are initializing BertModel from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).
- This IS NOT expected if you are initializing BertModel from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).

我担心部分模型权重未被使用或随机初始化会导致生成的嵌入不可信，想知道用PubMedBERT或BioBERT这类生物医学模型获取文本嵌入的正确方法是什么？

解决方案 #

1. 直接使用适配Sentence-Transformers的生物医学模型 #

优先选择HuggingFace上已经被包装为Sentence-Transformers格式的生物医学模型，这类模型已经做好池化层适配，不会出现权重未使用的警告，嵌入效果更稳定。常见选项包括：

• pritamdeka/BioBERT-mnli-snli-scinli-scitail-mednli-stsb：基于BioBERT微调的多任务适配模型
• gsarti/biobert-nli：适配自然语言推理任务的BioBERT Sentence-Transformers模型

使用示例：

from sentence_transformers import SentenceTransformer
import pandas as pd

# 加载适配好的模型
model = SentenceTransformer('pritamdeka/BioBERT-mnli-snli-scinli-scitail-mednli-stsb')
embeddings = model.encode(pd.Series(['hello', 'cell type', 'protein']))
print(embeddings)

2. 手动加载原生模型并自定义嵌入生成逻辑 #

如果一定要用原生PubMedBERT/BioBERT模型，可通过transformers库直接加载，自定义池化逻辑来生成嵌入，完全控制生成过程：

步骤1：加载原生模型与分词器

from transformers import BertTokenizer, BertModel
import torch
import pandas as pd

tokenizer = BertTokenizer.from_pretrained('microsoft/BiomedNLP-PubMedBERT-base-uncased-abstract-fulltext')
model = BertModel.from_pretrained('microsoft/BiomedNLP-PubMedBERT-base-uncased-abstract-fulltext')

步骤2：定义嵌入生成函数

支持两种常用池化方式：

• <[BOS_never_used_51bce0c785ca2f68081bfa7d91973934]>token池化：取模型输出中<[BOS_never_used_51bce0c785ca2f68081bfa7d91973934]>位置的隐藏状态作为句子嵌入
• 平均池化：对有效token（排除padding）的隐藏状态做平均

def get_embeddings(texts, tokenizer, model, device='cuda' if torch.cuda.is_available() else 'cpu'):
    model.to(device)
    model.eval()
    
    embeddings = []
    with torch.no_grad():
        for text in texts:
            inputs = tokenizer(text, return_tensors='pt', padding=True, truncation=True, max_length=512).to(device)
            outputs = model(**inputs)
            
            # 选项1：采用<[BOS_never_used_51bce0c785ca2f68081bfa7d91973934]>token池化
            cls_embedding = outputs.last_hidden_state[:, 0, :].cpu().numpy().squeeze()
            
            # 选项2：采用平均池化（排除padding token）
            # attention_mask = inputs['attention_mask']
            # token_embeddings = outputs.last_hidden_state
            # mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size())
            # sum_embeddings = torch.sum(token_embeddings * mask_expanded, 1)
            # sum_mask = torch.clamp(mask_expanded.sum(1), min=1e-9)
            # avg_embedding = (sum_embeddings / sum_mask).cpu().numpy().squeeze()
            
            embeddings.append(cls_embedding)
    
    return embeddings

# 生成文本嵌入
texts = pd.Series(['hello', 'cell type', 'protein'])
document_embeddings = get_embeddings(texts, tokenizer, model)
print(document_embeddings)

3. 正确看待警告，无需过度焦虑 #

你看到的未使用权重，是原模型预训练时用于**掩码语言建模（MLM）和下一句预测（NSP）**的任务头部权重，Sentence-Transformers加载时只会用到BERT的主体编码器，这些头部权重本来就不需要，因此这个警告是正常现象，不会导致嵌入不可信。但如果追求最优效果，还是建议使用前两种方法。

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