使用混合检索的Milvus向量存储¶
在这个笔记本中，我们将展示如何快速使用MilvusVectorStore进行混合检索的演示。（Milvus版本应该高于2.4.0）
如果您在colab上打开这个笔记本，您可能需要安装LlamaIndex 🦙。
%pip install llama-index-vector-stores-milvus
BGE-M3是FlagEmbedding中默认的稀疏嵌入方法，因此在安装llama-index时需要一并安装。
! pip install llama-index
! pip install FlagEmbedding
import logging
import sys
# 取消注释以查看调试日志
# logging.basicConfig(stream=sys.stdout, level=logging.DEBUG)
# logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))
from llama_index.core import VectorStoreIndex, SimpleDirectoryReader, Document
from llama_index.vector_stores.milvus import MilvusVectorStore
from IPython.display import Markdown, display
import textwrap
设置OpenAI¶
让我们首先添加OpenAI API密钥。这将允许我们访问OpenAI以获取嵌入和使用ChatGPT。
import openai
openai.api_key = "sk-"
下载数据¶
! mkdir -p 'data/paul_graham/'
! wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/docs/examples/data/paul_graham/paul_graham_essay.txt' -O 'data/paul_graham/paul_graham_essay.txt'
--2024-04-25 17:44:59-- https://raw.githubusercontent.com/run-llama/llama_index/main/docs/docs/examples/data/paul_graham/paul_graham_essay.txt Resolving raw.githubusercontent.com (raw.githubusercontent.com)... 185.199.111.133, 185.199.108.133, 185.199.109.133, ... Connecting to raw.githubusercontent.com (raw.githubusercontent.com)|185.199.111.133|:443... connected. HTTP request sent, awaiting response... 200 OK Length: 75042 (73K) [text/plain] Saving to: ‘data/paul_graham/paul_graham_essay.txt’ data/paul_graham/pa 100%[===================>] 73.28K --.-KB/s in 0.07s 2024-04-25 17:45:00 (994 KB/s) - ‘data/paul_graham/paul_graham_essay.txt’ saved [75042/75042]
生成我们的数据¶
有了我们的LLM集合，让我们开始使用Milvus索引。作为第一个例子，让我们从data/paul_graham/文件夹中的文件生成一个文档。在这个文件夹中，有一篇来自Paul Graham的单篇文章，标题为What I Worked On。为了生成这些文档，我们将使用SimpleDirectoryReader。
# 加载文档
documents = SimpleDirectoryReader("./data/paul_graham/").load_data()
print("文档ID:", documents[0].doc_id)
Document ID: ca3f5dbc-f772-41da-9a4f-bb4884691793
在数据中创建索引¶
现在我们有了一个文档，我们可以创建一个索引并插入文档。对于索引，我们将使用MilvusVectorStore。MilvusVectorStore接受一些参数：
uri (str, optional): 连接的URI，格式为"http://address:port"。默认为"http://localhost:19530"。token (str, optional): 登录的令牌。如果不使用rbac，则为空，如果使用rbac，则可能是"username:password"。默认为空字符串。collection_name (str, optional): 数据将被存储的集合的名称。默认为"llamalection"。dim (int, optional): 嵌入的维度。如果未提供，将在第一次插入时创建集合。默认为None。embedding_field (str, optional): 集合的嵌入字段的名称，默认为DEFAULT_EMBEDDING_KEY。doc_id_field (str, optional): 集合的doc_id字段的名称，默认为DEFAULT_DOC_ID_KEY。similarity_metric (str, optional): 要使用的相似度度量，目前支持IP和L2。默认为"IP"。consistency_level (str, optional): 用于新创建的集合的一致性级别。默认为"Strong"。overwrite (bool, optional): 是否覆盖同名的现有集合。默认为False。text_key (str, optional): 在传递的集合中存储文本的键。在带有自己的集合时使用。默认为None。index_config (dict, optional): 用于构建Milvus索引的配置。默认为None。search_config (dict, optional): 用于搜索Milvus索引的配置。注意，这必须与index_config指定的索引类型兼容。默认为None。batch_size (int): 在将数据插入Milvus时，配置在一个批处理中处理的文档数量。默认为DEFAULT_BATCH_SIZE。enable_sparse (bool): 一个布尔标志，指示是否启用对混合检索的稀疏嵌入的支持。默认为False。sparse_embedding_function (BaseSparseEmbeddingFunction, optional): 如果enable_sparse为True，则应提供此对象以将文本转换为稀疏嵌入。hybrid_ranker (str): 指定在混合搜索查询中使用的排名器类型。目前仅支持['RRFRanker'，'WeightedRanker']。默认为"RRFRanker"。hybrid_ranker_params (dict): 混合排名器的配置参数。对于"RRFRanker"，它应包括：
- 'k' (int): 用于Reciprocal Rank Fusion (RRF)的参数。该值用于计算排名分数，作为RRF算法的一部分，该算法将多个排名策略组合成单个分数，以提高搜索相关性。
对于"WeightedRanker"，它应包括：
- 'weights' (float列表): 由两个权重组成的列表：
- 密集嵌入组件的权重。
- 稀疏嵌入组件的权重。
这些权重用于调整嵌入的密集和稀疏组件在混合检索过程中的重要性。
- 'weights' (float列表): 由两个权重组成的列表：
默认为空字典，表示排名器将使用其预定义的默认设置。
现在，让我们开始创建一个用于混合检索的MilvusVectorStore。我们需要将enable_sparse设置为True以启用稀疏嵌入生成，还需要配置RRFRanker进行重新排序。更多详情，请参考Milvus重新排序。
# 在文档上创建索引
from llama_index.core import StorageContext
import os
vector_store = MilvusVectorStore(
dim=1536,
overwrite=True,
enable_sparse=True,
hybrid_ranker="RRFRanker",
hybrid_ranker_params={"k": 60},
)
storage_context = StorageContext.from_defaults(vector_store=vector_store)
index = VectorStoreIndex.from_documents(
documents, storage_context=storage_context
)
Sparse embedding function is not provided, using default.
Fetching 30 files: 0%| | 0/30 [00:00<?, ?it/s]
----------using 2*GPUs----------
查询数据¶
现在我们已经将文档存储在索引中，我们可以通过指定 vector_store_query_mode 来针对索引提出问题。索引将使用自身存储的数据作为chatgpt的知识库。
query_engine = index.as_query_engine(vector_store_query_mode="hybrid")
response = query_engine.query("What did the author learn?")
print(textwrap.fill(str(response), 100))
The author learned that the field of AI, as practiced at the time, was not as promising as initially believed. The author realized that the approach of using explicit data structures to represent concepts in AI was not effective in truly understanding natural language. This led the author to shift focus from traditional AI to exploring Lisp for its own merits, ultimately deciding to write a book about Lisp hacking.
response = query_engine.query("What was a hard moment for the author?")
print(textwrap.fill(str(response), 100))
Dealing with the stress and pressure related to managing Hacker News was a challenging moment for the author.
! pip install FlagEmbedding
from FlagEmbedding import BGEM3FlagModel
from typing import List
from llama_index.vector_stores.milvus.utils import BaseSparseEmbeddingFunction
class ExampleEmbeddingFunction(BaseSparseEmbeddingFunction):
def __init__(self):
self.model = BGEM3FlagModel("BAAI/bge-m3", use_fp16=False)
def encode_queries(self, queries: List[str]):
outputs = self.model.encode(
queries,
return_dense=False,
return_sparse=True,
return_colbert_vecs=False,
)["lexical_weights"]
return [self._to_standard_dict(output) for output in outputs]
def encode_documents(self, documents: List[str]):
outputs = self.model.encode(
documents,
return_dense=False,
return_sparse=True,
return_colbert_vecs=False,
)["lexical_weights"]
return [self._to_standard_dict(output) for output in outputs]
def _to_standard_dict(self, raw_output):
result = {}
for k in raw_output:
result[int(k)] = raw_output[k]
return result
现在我们可以在混合检索中使用这个。
vector_store = MilvusVectorStore(
dim=1536,
overwrite=True,
enable_sparse=True,
sparse_embedding_function=ExampleEmbeddingFunction(),
hybrid_ranker="RRFRanker",
hybrid_ranker_params={"k": 60},
)
Fetching 30 files: 0%| | 0/30 [00:00<?, ?it/s]
----------using 2*GPUs----------