BGE-M3模型结合Milvus向量数据库强强联合实现混合检索

在基于生成式人工智能的应用开发中，通过关键词或语义匹配的方式对用户提问意图进行识别是一个很重要的步骤，因为识别的精准与否会影响后续大语言模型能否检索出合适的内容作为推理的上下文信息（或选择合适的工具）以给出用户最符合预期的回答。

在本篇文章中，我将尽可能详细地介绍想达成准确识别用户提问意图的解决方案之一，即基于功能强大的BGE-M3模型和Milvus向量数据库实现混合检索（稠密向量 dense vector 和 稀疏向量 sparse vector）。

下面会先对混合检索，BGE-M3以及Milvus向量数据库做一个介绍，后面再对BGE-M3官方提供的使用Milvus的示例代码进行讲解，并基于自定义数据进行混合检索的实验效果进行呈现和分析。

何为混合检索

其实就是将多种单一检索技术结合起来使用的一种检索方式，目的是充分利用不同检索技术各自的优势，以达到更好的检索效果。接下来介绍两种比较常用的检索技术，嵌入检索（稠密检索）和 以BM25算法为代表的稀疏检索。

嵌入检索 (embedding retrieval)

指的是使用稠密向量（dense vector）表示来进行信息检索的方法。

稠密向量有以下特点：

1. 向量中大部分元素非零，是高维的实数向量，通常由深度学习模型生成；
2. 重点在于向量空间中的相似性计算，一般使用余弦相似度、点积等计算方法。这允许语义上相似的文本可以被较好理解，即使它们的词语重叠较少；
3. 能够捕捉文本的语义信息，而不是仅仅匹配关键词，但计算成本较高，需要更强大的计算资源。

稀疏检索 (sparse retrieval)

指的是使用稀疏向量（sparse vector）表示来进行信息检索的方法。

稀疏向量有以下特点：

1. 通常由词袋模型 (Bag-of-Words, BoW)、TF-IDF 等传统方法生成，向量维度通常对应词汇表中的词语，其值表示词频或 TF-IDF 权重，向量中大部分元素为零；
2. 主要基于关键词匹配，使用余弦相似度等方法计算向量之间的相似性；
3. 计算效率高，对语义理解能力有限。

区别总结（Gemini 1.5 Flash回答）：

总之，嵌入检索和稀疏检索是两种不同的信息检索方法，它们各有优缺点，实际应用中常常结合使用，形成混合检索系统，以达到更好的检索效果。

BGE-M3 模型

BAAI/bge-m3 是一个由北京人工智能研究院（BAAI）开源的嵌入模型，主要用于生成文本嵌入（text embeddings）。它是一个多功能、多语言、多粒度的模型。M3指的就是多功能性（Multi-Functionality）、多语言性（Multi-Linguality）和多粒度（Multi-Granularity）。

1. 多功能性 (Multi-Functionality): 该模型可以同时执行嵌入式模型的三个常见检索功能：稠密检索、多向量检索和稀疏检索。这意味着它能够灵活地应对不同的检索需求，并结合不同检索方法的优势。

2. 多语言性 (Multi-Linguality): 该模型支持100多种语言。 这意味着它可以处理多种语言的文本，并进行跨语言检索。

3. 多粒度性 (Multi-Granularity): 该模型能够处理不同粒度的输入，从短句到长达8192个词元的长文档。 这意味着它可以处理各种长度的文本，而不会受到输入长度的限制。

Milvus开源向量数据库介绍

官网介绍：https://milvus.io/intro

Github: https://github.com/milvus-io/milvus

Milvus 是一个以高效检索和高扩展性为特点的开源向量数据库，支持对大量的非结构化数据（如文本，图像还有多模态数据信息等）进行组织和检索。

Milvus 使用 Go 和 C++ 编程语言开发实现， 并通过CPU/GPU指令级优化，以实现最佳的向量搜索性能。

Milvus 提供多种本地部署方式，以适应不同规模和需求的应用场景（引导Gemini 1.5 Flash回答如下）：

1. 基于 Kubernetes (K8s) 的全分布式架构:

• 适用场景: 处理数万甚至数十万每秒查询请求，向量数据规模达到数十亿级别的大规模应用。需要高可扩展性和实时数据更新能力。
• 关键特性: 水平扩展性强，高可用性，实时数据处理。
• 主要差异: 部署最复杂，但性能和可扩展性最高。需要预先配置好 Kubernetes 集群。

安装部署参考：官方文档

2. 单机模式 (Docker或者Docker Compose):

• 适用场景: 小型应用、测试、开发和概念验证项目。相比全分布式模式，部署更简单。
• 关键特性: 使用 Docker 简化部署。适用于单节点环境。
• 主要差异: 可扩展性比分布式模式低。不适合高吞吐量或大型数据集。

安装部署参考：官方文档

3. Milvus Lite:

• 适用场景: 快速入门、学习和小型实验项目，尤其是在 Python 环境中。强调易用性和极简的设置。
• 关键特性: 通过 pip install 命令轻松安装。轻量级，快速上手。
• 主要差异: 可扩展性和功能最受限制。主要用于学习和原型设计，不适合生产环境部署或大型数据集。

安装部署参考：官方文档

总结表格:

简而言之，选择哪种部署模式取决于应用的规模和需求。对于海量数据集和高查询量，基于 Kubernetes 的部署是必要的。对于小型需求，单机模式或 Milvus Lite 提供更简单快速的设置。

本地实验前准备

下载 BGE-M3

git lfs install
git clone https://huggingface.co/BAAI/bge-m3

模型大小 total 2.2G

以Docker Compose方式部署Milvus & 下载 pymilvus 和 FlagEmbedding

在后续结合bge-m3模型实现基于语义的混合检索实验中，我们使用docker compose安装部署Milvus的方式，并搭配 pymilvus 第三方python库进行使用。

官方文档：Run Milvus with Docker Compose

运行命令如下

# 安装
wget https://github.com/milvus-io/milvus/releases/download/v2.5.0-beta/milvus-standalone-docker-compose.yml -O docker-compose.yml
# 运行
sudo docker-compose up -d
# 终端显示
Creating milvus-etcd  ... done
Creating milvus-minio ... done
Creating milvus-standalone ... done

pymilvus和FlagEmbedding（此库方便我们加载bge模型）安装命令：

pip install -U pymilvus
pip install -U FlagEmbedding

数据准备

因为我后期想搭一个基于提问内容自动匹配合适的api去调用并将返回结果作为大语言模型推理上下文信息的问答机器人小应用，所以我准备了对不同api调用功能目进行描述的几条**中文**短文本作为检索知识库内容。

我在聚合数据平台申请了三个api，分别是：

1. 黄金数据：提供黄金品种查询服务，支持黄金现货、黄金期货等品种查询；提供黄金价格查询服务，支持查询最新价、开盘价、最高价、最低价等价格信息。
2. 股票数据：根据股票编号查询股票信息，包括涨跌额、股票名称、开盘价、收盘价、当前价格、最高最低价、竞买价等信息；支持沪深、香港、美国三个市场的股票数据查询，覆盖全球各地的股票信息。
3. 新闻头条：提供最新的新闻头条信息，包括国内、国际、体育、娱乐、科技等各类资讯，让用户获取及时的新闻报道；支持按照不同的新闻分类进行查询，用户可以选择感兴趣的领域，如体育、科技、娱乐等，获取相关的新闻内容。

根据api的功能描述，建立检索知识库列表如下：

docs = [
    "提供黄金品种查询服务，支持黄金现货、黄金期货等品种查询；提供黄金价格查询服务，支持查询最新价、开盘价、最高价、最低价等价格信息。",
    "根据股票编号查询股票信息，包括涨跌额、股票名称、开盘价、收盘价、当前价格、最高最低价、竞买价等信息；支持沪深、香港、美国三个市场的股票数据查询，覆盖全球各地的股票信息。",
    "提供最新的新闻头条信息，包括国内、国际、体育、娱乐、科技等各类资讯，让用户获取及时的新闻报道；支持按照不同的新闻分类进行查询，用户可以选择感兴趣的领域，如体育、科技、娱乐等，获取相关的新闻内容。"]

代码实现

Milvus team 提供了使用 bge-m3 模型生成稀疏&稠密向量搭配Milvus向量数据库实现混合检索的示例代码，本次实验将以该代码作为参考，在其基础上进行适当修改调整。

1. 导入相关库；（本次实验不试用reranker进行精筛所以 use_reranker 设置为 False）

# A demo showing hybrid semantic search with dense and sparse vectors using Milvus.
#
# You can optionally choose to use the BGE-M3 model to embed the text as dense
# and sparse vectors, or simply use random generated vectors as an example.
#
# You can also use the BGE CrossEncoder model to rerank the search results.
#
# Note that the sparse vector search feature is only available in Milvus 2.4.0 or
# higher version. Make sure you follow https://milvus.io/docs/install_standalone-docker.md
# to set up the latest version of Milvus in your local environment.

# To connect to Milvus server, you need the python client library called pymilvus.
# To use BGE-M3 model, you need to install the optional `model` module in pymilvus.
# You can get them by simply running the following commands:
#
# pip install pymilvus
# pip install pymilvus[model]

# If true, use BGE-M3 model to generate dense and sparse vectors.
# If false, use random numbers to compose dense and sparse vectors.
use_bge_m3 = True
# If true, the search result will be reranked using BGE CrossEncoder model.
use_reranker = False

# The overall steps are as follows:
# 1. embed the text as dense and sparse vectors
# 2. setup a Milvus collection to store the dense and sparse vectors
# 3. insert the data to Milvus
# 4. search and inspect the result!
import random
import string
import numpy as np

from FlagEmbedding import BGEM3FlagModel

from pymilvus import (
    utility,
    FieldSchema, CollectionSchema, DataType,
    Collection, AnnSearchRequest, RRFRanker, connections, WeightedRanker
)

2. 指定bge-m3模型路径并以fp32精度（全精度浮点数）加载模型；

model_dir = "bge-m3"
model = BGEM3FlagModel(model_dir) 

3. 定义检索文本内容列表和用户提问示例；

docs = [
    "提供黄金品种查询服务，支持黄金现货、黄金期货等品种查询；提供黄金价格查询服务，支持查询最新价、开盘价、最高价、最低价等价格信息。",
    "根据股票编号查询股票信息，包括涨跌额、股票名称、开盘价、收盘价、当前价格、最高最低价、竞买价等信息；支持沪深、香港、美国三个市场的股票数据查询，覆盖全球各地的股票信息。",
    "提供最新的新闻头条信息，包括国内、国际、体育、娱乐、科技等各类资讯，让用户获取及时的新闻报道；支持按照不同的新闻分类进行查询，用户可以选择感兴趣的领域，如体育、科技、娱乐等，获取相关的新闻内容。"
]
# 用户提问示例
query = "请问黄金品种AU99.99价格如何？"

4. 将检索文本列表编码为稀疏和稠密向量；

passage_embeddings = model.encode(
    docs, 
    return_dense=True, 
    return_sparse=True, 
    return_colbert_vecs=False)

将 passage_embeddings 变量打印出来看看

{'dense_vecs': array([[-0.05132515,  0.0243751 , -0.06081232, ..., -0.03952702,
         0.02319626, -0.03473093],
       [-0.07099452, -0.01123387, -0.05009558, ..., -0.00663427,
         0.05279079,  0.01525173],
       [-0.03272525, -0.00253467, -0.04763371, ..., -0.00828949,
         0.0147302 , -0.05601017]], dtype=float32), 'lexical_weights': [defaultdict(<class 'int'>, {'6': 0.20276982, '2212': 0.17947777, '78563': 0.27358395, '182114': 0.24308087, '115106': 0.25357044, '3367': 0.19385284, '4': 0.17282556, '7499': 0.17129068, '15895': 0.11740263, '24781': 0.07431586, '153133': 0.18931584, '844': 0.06803232, '74': 0.10644381, '16806': 0.14869832, '16677': 0.054120254, '14498': 0.116243124, '4185': 0.07396346, '27964': 0.13261327, '37': 0.056129213, '14812': 0.059598763, '36315': 0.082842745, '5412': 0.09079999, '30': 0.011220258}), defaultdict(<class 'int'>, {'80476': 0.13337892, '50653': 0.23055996, '30862': 0.24856994, '5730': 0.24488933, '115106': 0.2854578, '5412': 0.22158799, '4': 0.17879836, '5384': 0.09052652, '85266': 0.091248944, '44040': 0.12723388, '22878': 0.106443286, '37': 0.06931204, '108279': 0.10546562, '4185': 0.05373253, '27964': 0.12790793, '14498': 0.093488105, '6633': 0.05225729, '36563': 0.106637016, '16806': 0.10283494, '14812': 0.045356337, '36315': 0.081632964, '97954': 0.14599079, '11795': 0.08332152, '844': 0.018887743, '74': 0.1779946, '7499': 0.18943872, '206056': 0.08312804, '6728': 0.12673198, '6665': 0.11565392, '3893': 0.115120836, '41953': 0.10232154, '60114': 0.1253739, '12833': 0.22091916, '78868': 0.14010657, '7153': 0.11074717, '40727': 0.100175254, '43': 0.08830682, '30': 0.025547896}), defaultdict(<class 'int'>, {'6': 0.21939242, '2212': 0.21460007, '110004': 0.16313581, '21606': 0.2181344, '3878': 0.15887201, '7781': 0.21819039, '5412': 0.19437945, '4': 0.15980814, '5384': 0.13158405, '13853': 0.108550005, '5559': 0.08813208, '28788': 0.17125762, '23942': 0.15054022, '7355': 0.15029378, '844': 0.0729285, '62335': 0.095689975, '144343': 0.23023936, '3933': 0.0954708, '10501': 0.16799492, '95579': 0.13417137, '32292': 0.12720525, '27835': 0.17556949, '74': 0.18231188, '7499': 0.20008457, '17236': 0.10246284, '13647': 0.12076213, '136229': 0.2379455, '115106': 0.23003161, '1441': 0.04709647, '8827': 0.11257933, '123701': 0.1811453, '17010': 0.15106374, '3322': 0.003149774, '84492': 0.07179478, '8105': 0.079666495, '30': 0.075698406})], 'colbert_vecs': None}

passage_embeddings 变量是一个dict数据结构，有三个key值，分别是 dense_vecs, lexical_weights, colbert_vecs。

passage_embeddings.keys()
# dict_keys(['dense_vecs', 'lexical_weights', 'colbert_vecs'])
passage_embeddings["dense_vecs"].shape
# (3, 1024) 3x1024维的一个dense vector
passage_embeddings["lexical_weights"]
# 生成的稀疏向量，sparse vector，存在一个list当中，有三个元素，元素类型是 collections.defaultdict
passage_embeddings["colbert_vecs"]
# colbert_vecs 这个Multi-Vector没生成，所以对应的值为 None

5. 将用户提问 query 编码为稀疏和稠密向量；

query_embeddings_raw = model.encode(
  [query], 
  return_dense=True, 
  return_sparse=True, 
  return_colbert_vecs=False)

6. 连接 Milvus（确保docker compose是up的状态）并创建 collection和index；

connections.connect("default", host="localhost", port="19530")
# 列出当前所有collection
utility.list_collections()
# 输出显示为空列表

基于schema创建一个新的名为 hybrid_demo 的collection

# Specify the data schema for the new Collection.
fields = [
    # Use auto generated id as primary key
    FieldSchema(name="pk", dtype=DataType.VARCHAR,
                is_primary=True, auto_id=True, max_length=100),
    # Store the original text to retrieve based on semantically distance
    FieldSchema(name="text", dtype=DataType.VARCHAR, max_length=512),
    # Milvus now supports both sparse and dense vectors, we can store each in
    # a separate field to conduct hybrid search on both vectors.
    FieldSchema(name="sparse_vector", dtype=DataType.SPARSE_FLOAT_VECTOR),
    # 注意这个的dim要和dense vector的维度保持一致，即1024
  	# 注意dtype也是要和dense vector的精度保持一致，即fp32
    FieldSchema(name="dense_vector", dtype=DataType.FLOAT_VECTOR,
                dim=1024),
]
schema = CollectionSchema(fields, "")
col_name = 'hybrid_demo'
# Now we can create the new collection with above name and schema.
col = Collection(col_name, schema, consistency_level="Strong")

接着我们需要为每一个向量字段创建index，创建的indices会被加载到内存中以实现高效检索。

# We need to create indices for the vector fields. The indices will be loaded into memory for efficient search.
sparse_index = {"index_type": "SPARSE_INVERTED_INDEX", "metric_type": "IP"}
col.create_index("sparse_vector", sparse_index)
dense_index = {"index_type": "FLAT", "metric_type": "IP"}
col.create_index("dense_vector", dense_index)
col.load()

7. 将生成的表示检索文本内容的稀疏&稠密向量加入到创建的collection中；

docs_embeddings = {}
docs_embeddings["dense"] = passage_embeddings["dense_vecs"]
docs_embeddings["sparse"] = passage_embeddings["lexical_weights"]
# 3. insert text and sparse/dense vector representations into the collection
entities = [docs, docs_embeddings["sparse"], docs_embeddings["dense"]]
col.insert(entities)
col.flush()

8. 开始混合检索

query_embeddings = {}
query_embeddings["dense"] = query_embeddings_raw["dense_vecs"]
query_embeddings["sparse"] = [dict(query_embeddings_raw["lexical_weights"][0])]

因为检索文档数量较少只有三条，所以top k中k的值设置为1

# 4. search and inspect the result!
k = 1 # we want to get the top 1 docs closest to the query

# Prepare the search requests for both vector fields
sparse_search_params = {"metric_type": "IP"}
sparse_req = AnnSearchRequest(query_embeddings["sparse"],
                              "sparse_vector", sparse_search_params, limit=k)
dense_search_params = {"metric_type": "IP"}
dense_req = AnnSearchRequest(query_embeddings["dense"],
                             "dense_vector", dense_search_params, limit=k)

# Search topK docs based on dense and sparse vectors and rerank with RRF.
res = col.hybrid_search([sparse_req, dense_req], rerank=RRFRanker(),
                        limit=k, output_fields=['text'])

# Currently Milvus only support 1 query in the same hybrid search request, so
# we inspect res[0] directly. In future release Milvus will accept batch
# hybrid search queries in the same call.
res = res[0]

打印出最为匹配的文档内容，检索正确：

print(res)
["id: 454220910858869856, distance: 0.032786883413791656, entity: {'text': '提供黄金品种查询服务，支持黄金现货、黄金期货等品种查询；提供黄金价格查询服务，支持查询最新价、开盘价、最高价、最低价等价格信息。'}"]

再试两个query：

Query 一：

query = "请问科技圈发生了什么热点事件？"
print(res)
["id: 454220910858869858, distance: 0.032786883413791656, entity: {'text': '提供最新的新闻头条信息，包括国内、国际、体育、娱乐、科技等各类资讯，让用户获取及时的新闻报道；支持按照不同的新闻分类进行查询，用户可以选择感兴趣的领域，如体育、科技、娱乐等，获取相关的新闻内容。'}"]

Query 二：

query = "请问股票代码为CSCO的美股今日走势如何？"
print(res)
["id: 454220910858869857, distance: 0.032786883413791656, entity: {'text': '根据股票编号查询股票信息，包括涨跌额、股票名称、开盘价、收盘价、当前价格、最高最低价、竞买价等信息；支持沪深、香港、美国三个市场的股票数据查询，覆盖全球各地的股票信息。'}"]

效果还是不错的。