使用BGE-M3和K近邻实现语义搜索

使用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 BGEM3FlagModel

model = 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]]