使用Qdrant+FastText实现向量存储和检索

news2024/11/17 1:30:51

1 概述

在《使用FastText库训练词向量》一文中，已经训练并保存好了一个用 FastText 训练的词向量模型-model_0809_001。在接下来的实践中，将直接使用该词向量模型来生成对应数据的向量，然后将向量和对应的负载存储到 Qdrant 中，然后进行向量检索。

2 向量存储和检索

下面是使用 Qdrant+FastText 实现向量存储和检索的实践案例。

2.1 创建 Qdrant 向量数据库

引入 qdrant_client 版本为 1.10.1 的 Qdrant 的 python 客户端库。

pip install qdrant-client
或
conda install conda-forge::qdrant-client

创建并启动一个Qdrant数据库，数据持久化到磁盘，数据存储地址 “/Users/Shared/data/qdrant”。

from qdrant_client import QdrantClient
 
client = QdrantClient(path="/Users/Shared/data/qdrant")

2.2 创建数据集

创建一个名称为 test_collection_0812 的数据集。其中向量维度为200维（和FastText词向量模型设置的向量维度保持一致），相似度计算方式使用欧几里德距离-Distance.EUCLID。

def create_collection_0812():
    client.create_collection(
        collection_name="test_collection_0812",
        vectors_config=VectorParams(size=200, distance=Distance.EUCLID),
    )

2.3 新增向量数据

首先使用 FastText 加载训练好的向量模型-model_0809_001，然后使用该模型生成特征数据的向量，并往 test_collection_0812 数据集中新增向量数据。

# 加载词向量模型
model = fasttext.load_model(model_0809_001)

def add_vectors_0812():
    # 往向量库新增数据
    client.upsert(
        collection_name="test_collection_0812",
        wait=True,
        points=[
            PointStruct(id=1, vector=model['黄疸茵陈片'], payload={"genericName": "黄疸茵陈片"}),
            PointStruct(id=2, vector=model['美沙拉嗪肠溶片'], payload={"genericName": "美沙拉嗪肠溶片"}),
            PointStruct(id=3, vector=model['盐酸贝那普利'], payload={"genericName": "盐酸贝那普利"}),
            PointStruct(id=4, vector=model['痔康片'], payload={"genericName": "痔康片"}),
            PointStruct(id=5, vector=model['阿普唑仑'], payload={"genericName": "阿普唑仑"}),
            PointStruct(id=6, vector=model['黄疸茵陈冲剂'], payload={"genericName": "黄疸茵陈冲剂"}),
            PointStruct(id=7, vector=model['肝素钙注射液'], payload={"genericName": "肝素钙注射液"}),
            PointStruct(id=8, vector=model['黄疸茵陈颗粒'], payload={"genericName": "黄疸茵陈颗粒"}),
        ],
    )

2.4 向量检索

进行向量检索时，同样是使用向量模型-model_0809_001生成查询入参的向量，然后进行向量搜索。搜索结果的分数（score）越小代表搜索结果与入参越相近。下面是搜索与“黄疸茵陈片”最相近的6条数据的案例。

def query_0812(genericName, size):
    search_result = client.search(
        collection_name="test_collection_0812", query_vector=model[genericName], limit=size, with_vectors=True
    )

    for var in search_result:
        print(var)


if __name__ == '__main__':
    query_0811("黄疸茵陈片", 6)

搜索结果如下所示：