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

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

# 文件大小约为700MB，因此需要一些时间来完成。
response = requests.get(embeddings_url, verify=True)  # 设置 `verify=False` 以绕过 SSL 验证
with open('vector_database_wikipedia_articles_embedded.zip', 'wb') as file:
    file.write(response.content)

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()

	id	url	title	text	title_vector	content_vector	vector_id
0	1	https://simple.wikipedia.org/wiki/April	April	April is the fourth month of the year in the J...	[0.001009464613161981, -0.020700545981526375, ...	[-0.011253940872848034, -0.013491976074874401,...	0
1	2	https://simple.wikipedia.org/wiki/August	August	August (Aug.) is the eighth month of the year ...	[0.0009286514250561595, 0.000820168002974242, ...	[0.0003609954728744924, 0.007262262050062418, ...	1
2	6	https://simple.wikipedia.org/wiki/Art	Art	Art is a creative activity that expresses imag...	[0.003393713850528002, 0.0061537534929811954, ...	[-0.004959689453244209, 0.015772193670272827, ...	2
3	8	https://simple.wikipedia.org/wiki/A	A	A or a is the first letter of the English alph...	[0.0153952119871974, -0.013759135268628597, 0....	[0.024894846603274345, -0.022186409682035446, ...	3
4	9	https://simple.wikipedia.org/wiki/Air	Air	Air refers to the Earth's atmosphere. Air is a...	[0.02224554680287838, -0.02044147066771984, -0...	[0.021524671465158463, 0.018522677943110466, -...	4

# 从字符串中读取向量并将其转换为列表
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只需要一个容器，但是在这个存储库的`./qdrant/docker-compose.yaml`中提供了docker-compose.yaml文件的示例。

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

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

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

qdrant.get_collections()

CollectionsResponse(collections=[])

索引数据

Qdrant将数据存储在__集合__中，其中每个对象至少由一个向量描述，并且可能包含一个名为__payload__的附加元数据。我们的集合将被称为文章，每个对象将由标题和内容向量描述。

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

from qdrant_client.http import models as rest

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

# Set up the collection with the vector configuration. You need to declare the vector size and distance metric for the collection. Distance metric enables vector database to index and search vectors efficiently.
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以存储向量和有效载荷
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 [01:34<00:00, 264.52it/s]

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

CountResult(count=25000)

搜索数据

一旦数据被放入Qdrant中，我们将开始查询集合中最接近的向量。我们可以提供一个额外的参数vector_name，以从基于标题切换到基于内容的搜索。

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.842)
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)

# This time we'll query using content vector
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)