如何基于BERT实现文本嵌入？求类似OpenAIEmbeddings()的函数

一、类似OpenAIEmbeddings()的封装实现 #

你可以通过两种方式快速实现BERT的文本嵌入功能，匹配OpenAIEmbeddings的使用体验：

1. 基于Transformers库手动封装

Hugging Face Transformers中没有直接的“BERTEmbeddings”类，但可以用BertModel和BertTokenizer快速封装：

from transformers import BertTokenizer, BertModel
import torch

class BertEmbeddings:
    def __init__(self, model_name="bert-base-uncased"):
        self.tokenizer = BertTokenizer.from_pretrained(model_name)
        self.model = BertModel.from_pretrained(model_name)
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.model.to(self.device)

    def embed_documents(self, texts):
        # 批量处理文本
        inputs = self.tokenizer(
            texts,
            padding=True,
            truncation=True,
            return_tensors="pt"
        ).to(self.device)
        
        with torch.no_grad():
            outputs = self.model(**inputs)
        
        # 可选两种嵌入方式：
        # 方式1：用<[BOS_never_used_51bce0c785ca2f68081bfa7d91973934]>token的预训练池化输出
        embeddings = outputs.pooler_output.cpu().numpy()
        # 方式2：取最后一层隐藏状态的均值（更适配句子嵌入任务）
        # embeddings = torch.mean(outputs.last_hidden_state, dim=1).cpu().numpy()
        
        return embeddings

# 使用示例
bert_embeddings = BertEmbeddings()
texts = ["Hello world", "BERT text embedding"]
embeds = bert_embeddings.embed_documents(texts)

2. 使用LangChain的HuggingFaceEmbeddings类

LangChain提供了和OpenAIEmbeddings风格一致的封装类，直接支持BERT系列模型：

from langchain.embeddings import HuggingFaceEmbeddings

# 初始化，指定BERT模型
bert_embeddings = HuggingFaceEmbeddings(model_name="bert-base-uncased")

# 调用方式与OpenAIEmbeddings完全对齐
texts = ["Hello world", "BERT text embedding"]
embeds = bert_embeddings.embed_documents(texts)
single_embed = bert_embeddings.embed_query("Single text query")

二、相关别称与官方资源说明 #

BERT的文本嵌入通常被称为上下文嵌入（Contextual Embeddings）或句子嵌入（Sentence Embeddings），官方文档中没有单独的“嵌入函数”章节，但在BertModel的输出说明里明确标注了用于生成嵌入的两个核心字段：

• last_hidden_state：模型最后一层的所有token隐藏状态，可通过均值/最大池化生成句子级嵌入
• pooler_output：经过预训练池化层处理后的<[BOS_never_used_51bce0c785ca2f68081bfa7d91973934]>token输出，适合快速生成句子表示

如果需要更优的语义嵌入效果，推荐使用专门微调过的BERT衍生模型（如sentence-transformers系列），这类模型在语义相似度、检索任务上表现更好，同样可以通过上述两种方式调用。

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