使用BGE-M3和K近邻实现语义搜索
1 简单介绍
随着大模型和 检索增强生成(Retrieval-augmented Generation, RAG)不断的发展,文本语义检索很重要。文本语义检索离不开文本向量生成,在文本向量生成上面有几个比较好的模型,例如:北京智源人工智能研究院构建的BGE(BAAI General Embedding)、北京希瑞亚斯科技有限公司构建的M3E(Moka Massive Mixed Embedding)、SentenceTransformers可以实现文本的句子、段落的向量化,也可实现图片的向量化。
本文选择BGE-M3模型,目前从公布的参数上看它的效果最好,参考地址如下:
# GitHub地址
https://github.com/FlagOpen/FlagEmbedding# HuggingFace地址
https://hf-mirror.com/BAAI/bge-m3
2 安装环境
# 安装BGE-M3
pip install -U FlagEmbedding -i https://pypi.tuna.tsinghua.edu.cn/simple# 安装sklean
pip install scikit-learn -i https://pypi.tuna.tsinghua.edu.cn/simple
3 简单使用BGE
3.1 使用BGE生成文本向量
import numpy as np
from FlagEmbedding import BGEM3FlagModel# 引入模型
model = BGEM3FlagModel(model_name_or_path="E:/model/bge-m3")# 注意:单篇文章的字数最大8192个
sentences = ["文章摘要1","文章摘要2"
]# 对文本进行向量编码,使用稠密编码
# sentences_vector_dict = model.encode(sentences, return_dense=True)
sentences_vector_dict = model.encode(sentences,# 设置返回Dense Embedding,默认打开return_dense=True,# 设置返回Sparse Embedding (Lexical Weight),默认关闭return_sparse=True,# 设置返回Multi-Vector (ColBERT),默认关闭return_colbert_vecs=True
)
#
"""
返回一个对象,对象中含有三个属性分别是:dense_vecs、lexical_weights、colbert_vecs
{"dense_vecs": array([[……], [……]],"lexical_weights": [dict({'6': 0.002955027, '20403': 0.107134104, '89095': 0.25817403, '418': 0.075165465}), dict({'20403': 0.08092275, '89095': 0.22856326, '304': 0.079813644})],"colbert_vecs": array([[……], [……]],
}
"""
print(sentences_vector_dict)# 获取稠密编码
sentences_vector_arr = sentences_vector_dict.get("dense_vecs")# 查看类型:numpy.ndarray
print(type(sentences_vector_arr))# 查看维度:(2, 1024)
print(np.shape(sentences_vector_arr))# 获取第一个文本的语义向量:一维向量[……]
print(sentences_vector_arr[0])# 使用sentence_transformers比较两个文本的相似度
from sentence_transformers import util
print(util.cos_sim(sentences_vector_arr[0], sentences_vector_arr[1]))
3.2 使用BGE计算colbert的评分
from FlagEmbedding import BGEM3FlagModelmodel = BGEM3FlagModel(model_name_or_path="E:/model/bge-m3", use_fp16=True)sentences_1 = ["What is BGE M3?", "Defination of BM25"]
sentences_2 = ["BGE M3 is an embedding model supporting dense retrieval, lexical matching and multi-vector interaction.","BM25 is a bag-of-words retrieval function that ranks a set of documents based on the query terms appearing in each document"]output_1 = model.encode(sentences_1, return_dense=True, return_sparse=True, return_colbert_vecs=True)
output_2 = model.encode(sentences_2, return_dense=True, return_sparse=True, return_colbert_vecs=True)# colbert_score中使用了”爱因斯坦求和“(torch.einsum),colbert_score具体算法请参考相关文献
"""
爱因斯坦求和是一种对求和公式简洁高效的记法,其原则是当变量下标重复出现时,即可省略繁琐的求和符号。
例如:
C = einsum('ij,jk->ik', A, B)注释:->符号就相当于等号,->左边为输入,右边为输出。
"""
print(model.colbert_score(output_1['colbert_vecs'][0], output_2['colbert_vecs'][0]))
print(model.colbert_score(output_1['colbert_vecs'][0], output_2['colbert_vecs'][1]))
3.3 使用BGE-M3和K近邻实现语义检索
from FlagEmbedding import BGEM3FlagModel
from sklearn.neighbors import NearestNeighbors# 引入模型
model = BGEM3FlagModel(model_name_or_path="E:/model/bge-m3")# 注意:单篇文章的字数最大8192个
sentences = ["河南大学是一所历史悠久的综合性高校","河南大学软件学院","软件工程是一个专业","我是河南大学的一名学生"
]# 对文本进行向量编码,使用稠密编码
sentences_vector_dict = model.encode(sentences)# 获取稠密编码
sentences_vector_arr = sentences_vector_dict.get("dense_vecs")# 初始化最近邻搜索对象,n_neighbors:距离最近的数量,algorithm:使用的算法
nbrs = NearestNeighbors(n_neighbors=3, algorithm='ball_tree').fit(sentences_vector_arr)# 查询最近邻
distances, indices = nbrs.kneighbors(sentences_vector_arr)# 注意:距离的第一个值是0,因为文本自己到自己是最近的才为0;
print("距离:", distances)
# 索引中第一个值是文本的索引位置,对应自己,因为文本自己到自己最近
print("索引:", indices)
输出:
距离:
[[0. 0.69528755 0.8755811 ][0. 0.77477873 0.86843368][0. 0.86843368 1.01228235][0. 0.69528755 0.77477873]]
索引:
[[0 3 1][1 3 2][2 1 3][3 0 1]]