BCE embedding模型部署和推理

BCEmbedding是由网易有道开发的中英双语和跨语种语义表征算法模型库，其中包含 EmbeddingModel和 RerankerModel两类基础模型。EmbeddingModel专门用于生成语义向量，在语义搜索和问答中起着关键作用，而 RerankerModel擅长优化语义搜索结果和语义相关顺序精排。

模型部署推理

前期准备

创建一个conda环境 

conda create --name bce python=3.10  # 创建新环境
conda activate bce

官网安装PyTorch

也可以离线安装PyTorch  https://download.pytorch.org/whl/torch_stable.html

安装transformers和sentence-transformers

pip install transformers
pip install sentence-transformers

模型下载

安装modelscope

pip install modelscope

下载bce模型

modelscope download --model netease-youdao/bce-embedding-base_v1

默认情况下，modelscope 将模型存储在以下位置：

• Linux/macOS: ~/.cache/modelscope/
• Windows: C:\Users<用户名>.cache\modelscope

模型使用

有三种使用方式

1. 安装BCEmbedding

   pip install BCEmbedding==0.1.5

   或者使用源码安装

   git clone https://github.com/netease-youdao/BCEmbedding/tree/master
   cd BCEmbedding
   pip install -v -e .

2. 使用transformers库
3. 使用sentence_transformers库

# 方式一，直接使用官方封装好的库BCEmbedding
from BCEmbedding import EmbeddingModel
sentences = ['姚明', '姚明老婆']
device = 'cuda:2'  
# init embedding model
model = EmbeddingModel(model_name_or_path="/bce-embedding-base_v1", device=device)
# extract embeddings
embeddings = model.encode(sentences)
print(embeddings)
'''
[[ 0.00243835  0.03394481  0.02659898 ...  0.02342649 -0.04681893
  -0.07420494]
 [ 0.02900401  0.0521045   0.0524049  ... -0.00160577 -0.0363229
   0.00262558]]
'''

# 方式二，使用transformers加载模型
from transformers import AutoModel, AutoTokenizer
sentences = ['姚明', '姚明老婆']
# 加载模型
tokenizer = AutoTokenizer.from_pretrained('/bce-embedding-base_v1')
model = AutoModel.from_pretrained('/bce-embedding-base_v1')
device = 'cuda:2'  
model.to(device)
# tokenizer
inputs = tokenizer(sentences, padding=True, truncation=True, max_length=512, return_tensors="pt")
inputs_on_device = {k: v.to(device) for k, v in inputs.items()}
# embedding
outputs = model(**inputs_on_device, return_dict=True)
embeddings = outputs.last_hidden_state[:, 0]  
embeddings = embeddings / embeddings.norm(dim=1, keepdim=True)  # normalize
print(embeddings)
'''
tensor([[ 0.0024,  0.0339,  0.0266,  ...,  0.0234, -0.0468, -0.0742],
        [ 0.0290,  0.0521,  0.0524,  ..., -0.0016, -0.0363,  0.0026]],
       device='cuda:2', grad_fn=<DivBackward0>)
'''

# 方式三，使用sentence_transformers加载模型
from sentence_transformers import SentenceTransformer
sentences = ['姚明', '姚明老婆']
device = 'cuda:2'  
model = SentenceTransformer("/bce-embedding-base_v1", device=device)
# embeddings
embeddings = model.encode(sentences, normalize_embeddings=True)
print(embeddings)
'''
[[ 0.00243835  0.03394481  0.02659898 ...  0.02342649 -0.04681893
  -0.07420494]
 [ 0.02900401  0.05210449  0.0524049  ... -0.00160577 -0.0363229
   0.00262558]]
'''

关于这三种加载方式源码的分析

方式二，transformers加载模型为基础的加载方式，首先对输入文本进行tokenizer处理，然后输入到embedding模型中，得到的维度是[batch_size, sequence_length, hidden_size]

• 2 (batch_size): 输入了两个句子，分别是 '姚明' 和 '姚明老婆'，因此 batch_size = 2。
• 6 (sequence_length): 这是两个句子在经过 tokenizer 处理后，生成的 token 序列的长度。虽然这两个句子在原始形式中非常简短（分别有 2 个和 4 个字符），但 tokenizer 将它们转为 6 个 token，包括特殊的 [CLS] 和 [SEP] token，以及对中文字符的分词。
• 768 (hidden_size): 这是模型的隐藏层大小，每个 token 都会映射到一个 768 维的向量。

通过 outputs.last_hidden_state[:, 0] 获取每个句子第一个 token（[CLS] token）的向量。因为通常 [CLS] token 的向量表示整个句子的语义，获得了形状为 (2, 768) 的嵌入，即每个句子的表示向量

方式一，直接使用官方封装好的库BCEmbedding，内部原理和方式二相似，我们可以在源码中找到一样的代码。

inputs = self.tokenizer(
         sentence_batch, 
         padding=True,
         truncation=True,
         max_length=max_length,
         return_tensors="pt"
)
inputs_on_device = {k: v.to(self.device) for k, v in inputs.items()}
outputs = self.model(**inputs_on_device, return_dict=True)

if self.pooler == "cls":
    embeddings = outputs.last_hidden_state[:, 0]
    
if normalize_to_unit:
    embeddings = embeddings / embeddings.norm(dim=1, keepdim=True)
    embeddings_collection.append(embeddings.cpu())

不做过多赘述。

方式三，使用sentence_transformers加载模型

方式三的实现和另外两种方式不同

在源码中转换为token形式如下面的代码

# 将句子转换为模型可接受的 token 格式
features = self.tokenize(sentences_batch)

然后进行embedding操作

out_features = self.forward(features)

用了模型的前向传播 (forward) 函数，获取了所有 token 的嵌入，即 last_hidden_state

out_features["sentence_embedding"] = truncate_embeddings(
    out_features["sentence_embedding"], self.truncate_dim
)

sentence_embedding 被提取为每个句子的表示。这通常是通过提取 CLS token 的嵌入来实现的

embedding后计算不同句子的相似度

from sentence_transformers import SentenceTransformer
from sentence_transformers.util import cos_sim
sentences = ['姚明', '姚明老婆', '奥尼尔', '中国上海']
device = 'cuda:2' 
model = SentenceTransformer("/bce-embedding-base_v1", device=device)
embeddings = model.encode(sentences, normalize_embeddings=True)
sim = model.similarity(embeddings,embeddings)
print(sim)

'''
tensor([[1.0000, 0.7057, 0.5424, 0.3242],
        [0.7057, 1.0000, 0.4712, 0.3616],
        [0.5424, 0.4712, 1.0000, 0.2706],
        [0.3242, 0.3616, 0.2706, 1.0000]])
'''

BCEmbedding/README_zh.md at master · netease-youdao/BCEmbedding (github.com)

为RAG而生-BCE embedding技术报告 - 知乎 (zhihu.com)

BCEmbedding/README_zh.md at master · netease-youdao/BCEmbedding (github.com)

NLP：如何计算两个句子的相似度-CSDN博客

csdn这编辑器做的真不好用，从自己的笔记软件下复制到编辑器中，全是问题，强迫症看的难受。