使用 Qdrant 进行嵌入搜索

本笔记本将引导您完成一个简单的流程，以下载一些数据、对其进行嵌入，然后使用精选的向量数据库对其进行索引和搜索。这是客户希望在安全的环境中存储和搜索我们的嵌入以支持生产用例（如聊天机器人、主题建模等）的常见需求。

什么是向量数据库

向量数据库是一种用于存储、管理和搜索嵌入向量的数据库。近年来，由于人工智能在解决涉及自然语言、图像识别和其他非结构化数据形式的用例方面越来越有效，使用嵌入将非结构化数据（文本、音频、视频等）编码为向量以供机器学习模型使用的做法呈爆炸式增长。向量数据库已成为企业交付和扩展这些用例的有效解决方案。

为什么使用向量数据库

向量数据库使企业能够利用我们在此仓库中共享的许多嵌入用例（例如，问答、聊天机器人和推荐服务），并在安全、可扩展的环境中使用它们。我们的许多客户使用嵌入在小规模上解决他们的问题，但性能和安全性阻碍了他们投入生产——我们将向量数据库视为解决此问题的关键组成部分，在本指南中，我们将介绍嵌入文本数据、将其存储在向量数据库中以及使用它进行语义搜索的基础知识。

演示流程

演示流程如下：

• 设置：导入包并设置任何必需的变量
• 加载数据：加载数据集并使用 OpenAI 嵌入对其进行嵌入
• Qdrant
  • 设置：在这里我们将设置 Qdrant 的 Python 客户端。更多详情请参见此处
  • 索引数据：我们将创建一个包含 标题 和 内容 向量的集合
  • 搜索数据：我们将运行几次搜索以确认其有效

运行完此笔记本后，您应该对如何设置和使用向量数据库有一个基本的了解，然后可以继续进行更复杂的用例，利用我们的嵌入。

设置

导入所需的库并设置我们想要使用的嵌入模型。

# 我们需要安装 Qdrant 客户端
!pip install qdrant-client

import openai
import pandas as pd
from ast import literal_eval
import qdrant_client # Qdrant 的 Python 客户端库

# 这可以更改为您选择的嵌入模型。确保它与用于生成嵌入的模型相同
EMBEDDING_MODEL = "text-embedding-ada-002"

# 忽略未关闭的 SSL 套接字警告 - 如果遇到这些错误，可以选择忽略
import warnings

warnings.filterwarnings(action="ignore", message="unclosed", category=ResourceWarning)
warnings.filterwarnings("ignore", category=DeprecationWarning)

加载数据

在本节中，我们将加载在此会话之前准备好的嵌入数据。

import wget

embeddings_url = "https://cdn.openai.com/API/examples/data/vector_database_wikipedia_articles_embedded.zip"

# 该文件大约为 700 MB，因此需要一些时间
wget.download(embeddings_url)

100% [......................................................................] 698933052 / 698933052




'vector_database_wikipedia_articles_embedded (10).zip'

import zipfile
with zipfile.ZipFile("vector_database_wikipedia_articles_embedded.zip","r") as zip_ref:
    zip_ref.extractall("../data")

article_df = pd.read_csv('../data/vector_database_wikipedia_articles_embedded.csv')

article_df.head()

# 将向量从字符串读回列表
article_df['title_vector'] = article_df.title_vector.apply(literal_eval)
article_df['content_vector'] = article_df.content_vector.apply(literal_eval)

# 将 vector_id 设置为字符串
article_df['vector_id'] = article_df['vector_id'].apply(str)

article_df.info(show_counts=True)

<class 'pandas.core.frame.DataFrame'>
RangeIndex: 25000 entries, 0 to 24999
Data columns (total 7 columns):
 #   Column          Non-Null Count  Dtype
---  ------          --------------  -----
 0   id              25000 non-null  int64
 1   url             25000 non-null  object
 2   title           25000 non-null  object
 3   text            25000 non-null  object
 4   title_vector    25000 non-null  object
 5   content_vector  25000 non-null  object
 6   vector_id       25000 non-null  object
dtypes: int64(1), object(6)
memory usage: 1.3+ MB

Qdrant

Qdrant 是一个用 Rust 编写的高性能向量搜索数据库。它同时提供本地和云版本，但为了演示的目的，我们将使用本地部署模式。

设置所有内容将需要：

• 启动一个本地 Qdrant 实例
• 配置集合并将数据存储在其中
• 尝试进行一些查询

设置

对于本地部署，我们将使用 Docker，根据 Qdrant 文档：https://qdrant.tech/documentation/quick_start/。Qdrant 只需要一个容器，但 docker-compose.yaml 文件的示例在此仓库的 ./qdrant/docker-compose.yaml 中可用。

您可以通过导航到此目录并运行 docker-compose up -d 来在本地启动 Qdrant 实例

您可能需要将 Docker 的内存限制提高到 8GB 或更多。否则 Qdrant 可能会因类似 7 Killed 的错误消息而失败。

! docker compose up -d

 [1A [1B [0G [?25l[+] Running 1/0
  [32m✔ [0m Container qdrant-qdrant-1   [32mRunning [0m                                       [34m0.0s  [0m
 [?25h

qdrant = qdrant_client.QdrantClient(host="localhost", port=6333)

qdrant.get_collections()

CollectionsResponse(collections=[CollectionDescription(name='Articles')])

索引数据

Qdrant 将数据存储在 collections 中，其中每个对象至少由一个向量描述，并且可以包含称为 payload 的附加元数据。我们的集合将被称为 Articles，每个对象将由 title 和 content 向量描述。

我们将使用官方的 qdrant-client 包，其中已内置所有实用方法。

from qdrant_client.http import models as rest

# 获取第一行的向量大小以设置集合
vector_size = len(article_df['content_vector'][0])

# 使用向量配置设置集合。您需要为集合声明向量大小和距离度量。距离度量使向量数据库能够有效地索引和搜索向量。
qdrant.recreate_collection(
    collection_name='Articles',
    vectors_config={
        'title': rest.VectorParams(
            distance=rest.Distance.COSINE,
            size=vector_size,
        ),
        'content': rest.VectorParams(
            distance=rest.Distance.COSINE,
            size=vector_size,
        ),
    }
)

True

vector_size = len(article_df['content_vector'][0])

qdrant.recreate_collection(
    collection_name='Articles',
    vectors_config={
        'title': rest.VectorParams(
            distance=rest.Distance.COSINE,
            size=vector_size,
        ),
        'content': rest.VectorParams(
            distance=rest.Distance.COSINE,
            size=vector_size,
        ),
    }
)

True

除了在 vector 下定义的向量配置外，我们还可以定义 payload 配置。Payload 是一个可选字段，允许您将其他元数据与向量一起存储。在我们的例子中，我们将存储文章的 id、title 和 url。由于我们在搜索结果中从 payload 返回最近文章的标题，因此我们还可以为用户提供文章的 URL（它是元数据的一部分）。

from qdrant_client.models import PointStruct # 导入 PointStruct 以存储向量和 payload
from tqdm import tqdm # 用于显示进度条的库

# 使用 tqdm 填充集合向量以显示进度
for k, v in tqdm(article_df.iterrows(), desc="Upserting articles", total=len(article_df)):
    try:
        qdrant.upsert(
            collection_name='Articles',
            points=[
                PointStruct(
                    id=k,
                    vector={'title': v['title_vector'],
                            'content': v['content_vector']},
                    payload={
                        'id': v['id'],
                        'title': v['title'],
                        'url': v['url']
                    }
                )
            ]
        )
    except Exception as e:
        print(f"Failed to upsert row {k}: {v}")
        print(f"Exception: {e}")

Upserting articles: 100%|█████████████████████████████████████████████████████████████████████████████████████| 25000/25000 [02:52<00:00, 144.82it/s]

# 检查集合大小以确保所有点都已存储
qdrant.count(collection_name='Articles')

CountResult(count=25000)

搜索数据

将数据放入 Qdrant 后，我们将开始查询集合以查找最接近的向量。我们可以提供附加参数 vector_name 来切换基于标题或内容的搜索。请确保使用 text-embedding-ada-002 模型，因为文件中的原始嵌入是使用此模型创建的。

def query_qdrant(query, collection_name, vector_name='title', top_k=20):

    # 从用户查询创建嵌入向量
    embedded_query = openai.embeddings.create(
        input=query,
        model=EMBEDDING_MODEL,
    ).data[0].embedding # 我们取列表中的第一个嵌入

    query_results = qdrant.search(
        collection_name=collection_name,
        query_vector=(
            vector_name, embedded_query
        ),
        limit=top_k,
        query_filter=None
    )

    return query_results

query_results = query_qdrant('modern art in Europe', 'Articles', 'title')
for i, article in enumerate(query_results):
    print(f'{i + 1}. {article.payload["title"]}, URL: {article.payload["url"]} (Score: {round(article.score, 3)})')

1. Museum of Modern Art, URL: https://simple.wikipedia.org/wiki/Museum%20of%20Modern%20Art (Score: 0.875)
2. Western Europe, URL: https://simple.wikipedia.org/wiki/Western%20Europe (Score: 0.867)
3. Renaissance art, URL: https://simple.wikipedia.org/wiki/Renaissance%20art (Score: 0.864)
4. Pop art, URL: https://simple.wikipedia.org/wiki/Pop%20art (Score: 0.86)
5. Northern Europe, URL: https://simple.wikipedia.org/wiki/Northern%20Europe (Score: 0.855)
6. Hellenistic art, URL: https://simple.wikipedia.org/wiki/Hellenistic%20art (Score: 0.853)
7. Modernist literature, URL: https://simple.wikipedia.org/wiki/Modernist%20literature (Score: 0.847)
8. Art film, URL: https://simple.wikipedia.org/wiki/Art%20film (Score: 0.843)
9. Central Europe, URL: https://simple.wikipedia.org/wiki/Central%20Europe (Score: 0.843)
10. European, URL: https://simple.wikipedia.org/wiki/European (Score: 0.841)
11. Art, URL: https://simple.wikipedia.org/wiki/Art (Score: 0.841)
12. Byzantine art, URL: https://simple.wikipedia.org/wiki/Byzantine%20art (Score: 0.841)
13. Postmodernism, URL: https://simple.wikipedia.org/wiki/Postmodernism (Score: 0.84)
14. Eastern Europe, URL: https://simple.wikipedia.org/wiki/Eastern%20Europe (Score: 0.839)
15. Cubism, URL: https://simple.wikipedia.org/wiki/Cubism (Score: 0.839)
16. Europe, URL: https://simple.wikipedia.org/wiki/Europe (Score: 0.839)
17. Impressionism, URL: https://simple.wikipedia.org/wiki/Impressionism (Score: 0.838)
18. Bauhaus, URL: https://simple.wikipedia.org/wiki/Bauhaus (Score: 0.838)
19. Surrealism, URL: https://simple.wikipedia.org/wiki/Surrealism (Score: 0.837)
20. Expressionism, URL: https://simple.wikipedia.org/wiki/Expressionism (Score: 0.837)

# 这次我们将使用内容向量进行查询
query_results = query_qdrant('Famous battles in Scottish history', 'Articles', 'content')
for i, article in enumerate(query_results):
    print(f'{i + 1}. {article.payload["title"]}, URL: {article.payload["url"]} (Score: {round(article.score, 3)})')

1. Battle of Bannockburn, URL: https://simple.wikipedia.org/wiki/Battle%20of%20Bannockburn (Score: 0.869)
2. Wars of Scottish Independence, URL: https://simple.wikipedia.org/wiki/Wars%20of%20Scottish%20Independence (Score: 0.861)
3. 1651, URL: https://simple.wikipedia.org/wiki/1651 (Score: 0.852)
4. First War of Scottish Independence, URL: https://simple.wikipedia.org/wiki/First%20War%20of%20Scottish%20Independence (Score: 0.85)
5. Robert I of Scotland, URL: https://simple.wikipedia.org/wiki/Robert%20I%20of%20Scotland (Score: 0.846)
6. 841, URL: https://simple.wikipedia.org/wiki/841 (Score: 0.844)
7. 1716, URL: https://simple.wikipedia.org/wiki/1716 (Score: 0.844)
8. 1314, URL: https://simple.wikipedia.org/wiki/1314 (Score: 0.837)
9. 1263, URL: https://simple.wikipedia.org/wiki/1263 (Score: 0.836)
10. William Wallace, URL: https://simple.wikipedia.org/wiki/William%20Wallace (Score: 0.835)
11. Stirling, URL: https://simple.wikipedia.org/wiki/Stirling (Score: 0.831)
12. 1306, URL: https://simple.wikipedia.org/wiki/1306 (Score: 0.831)
13. 1746, URL: https://simple.wikipedia.org/wiki/1746 (Score: 0.83)
14. 1040s, URL: https://simple.wikipedia.org/wiki/1040s (Score: 0.828)
15. 1106, URL: https://simple.wikipedia.org/wiki/1106 (Score: 0.827)
16. 1304, URL: https://simple.wikipedia.org/wiki/1304 (Score: 0.826)
17. David II of Scotland, URL: https://simple.wikipedia.org/wiki/David%20II%20of%20Scotland (Score: 0.825)
18. Braveheart, URL: https://simple.wikipedia.org/wiki/Braveheart (Score: 0.824)
19. 1124, URL: https://simple.wikipedia.org/wiki/1124 (Score: 0.824)
20. July 27, URL: https://simple.wikipedia.org/wiki/July%2027 (Score: 0.823)