使用FastEmbed和Qdrant构建混合搜索服务

时间: 20 分钟级别: 初学者输出: GitHub

本教程向您展示如何构建和部署您自己的混合搜索服务，以浏览来自startups-list.com的公司描述，并挑选与您的查询最相似的公司。 该网站包含公司名称、描述、位置和每个条目的图片。

正如我们已经在博客上写过的，混合搜索没有一个单一的定义。 在本教程中，我们涵盖了密集和稀疏嵌入结合的情况。 前者指的是由BERT等知名神经网络生成的嵌入，而后者则更接近于传统的全文搜索方法。

我们的混合搜索服务将使用快速嵌入包来生成文本描述的嵌入，并使用FastAPI来提供搜索API。 Fastembed原生与Qdrant客户端集成，因此您可以轻松地将数据上传到Qdrant并执行搜索查询。

Hybrid Search Schema

工作流程

要创建一个混合搜索服务，您需要先转换原始数据，然后创建一个搜索函数来操作它。 首先，您将1）使用修改版的BERT ML模型下载并准备一个样本数据集。接着，您将2）将数据加载到Qdrant中，3）创建一个混合搜索API，并4）使用FastAPI提供服务。

Hybrid Search Workflow

先决条件

要完成本教程，您将需要：

  • Docker - 使用 Qdrant 的最简单方法是运行预构建的 Docker 镜像。
  • 原始解析数据 来自 startups-list.com。
  • Python 版本 >=3.8

准备样本数据集

要对创业公司描述进行混合搜索，首先必须将描述数据编码为向量。 将Fastembed集成到qdrant客户端中，将编码和上传合并为一个步骤。

它还负责批处理和并行化，因此您不必担心。

让我们从下载数据和安装必要的包开始。

  1. First you need to download the dataset.
wget https://storage.googleapis.com/generall-shared-data/startups_demo.json

在 Docker 中运行 Qdrant

接下来，您需要使用向量引擎管理所有数据。Qdrant 允许您存储、更新或删除创建的向量。最重要的是，它允许您通过方便的 API 搜索最近的向量。

注意： 在开始之前，请创建一个项目目录并在其中创建一个虚拟的python环境。

  1. Download the Qdrant image from DockerHub.
docker pull qdrant/qdrant
  1. Start Qdrant inside of Docker.
docker run -p 6333:6333 \
    -v $(pwd)/qdrant_storage:/qdrant/storage \
    qdrant/qdrant

你应该看到像这样的输出

...
[2021-02-05T00:08:51Z INFO  actix_server::builder] Starting 12 workers
[2021-02-05T00:08:51Z INFO  actix_server::builder] Starting "actix-web-service-0.0.0.0:6333" service on 0.0.0.0:6333

通过访问http://localhost:6333/来测试服务。你应该能在浏览器中看到Qdrant的版本信息。

所有上传到Qdrant的数据都保存在./qdrant_storage目录中，即使你重新创建容器，数据也会被持久化。

上传数据到Qdrant

  1. Install the official Python client to best interact with Qdrant.
pip install "qdrant-client[fastembed]>=1.8.2"

注意： 本教程需要 fastembed 版本 >=0.2.6。

此时，您应该在startups_demo.json文件中拥有启动记录，并且Qdrant在本地机器上运行。

现在你需要编写一个脚本，将所有启动数据和向量上传到搜索引擎中。

  1. Create a client object for Qdrant.
# Import client library
from qdrant_client import QdrantClient

client = QdrantClient(url="http://localhost:6333")
  1. Select model to encode your data.

你将使用两个预训练模型来分别计算密集和稀疏向量：sentence-transformers/all-MiniLM-L6-v2 和 prithivida/Splade_PP_en_v1。

client.set_model("sentence-transformers/all-MiniLM-L6-v2")
# comment this line to use dense vectors only
client.set_sparse_model("prithivida/Splade_PP_en_v1")
  1. Related vectors need to be added to a collection. Create a new collection for your startup vectors.
if not client.collection_exists("startups"):
    client.create_collection(
        collection_name="startups",
        vectors_config=client.get_fastembed_vector_params(),
        # comment this line to use dense vectors only
        sparse_vectors_config=client.get_fastembed_sparse_vector_params(),  
    )

Qdrant 要求向量具有自己的名称和配置。

方法 get_fastembed_vector_params 和 get_fastembed_sparse_vector_params 帮助您获取您正在使用的模型的相应参数。 这些参数包括向量大小、距离函数等。

如果没有fastembed集成，您需要手动指定向量大小和距离函数。了解更多信息这里。

此外，您可以为您的向量指定扩展配置，例如 quantization_config 或 hnsw_config。

  1. Read data from the file.
import json

payload_path = "startups_demo.json"
metadata = []
documents = []

with open(payload_path) as fd:
    for line in fd:
        obj = json.loads(line)
        documents.append(obj.pop("description"))
        metadata.append(obj)

在这段代码中，我们从startups_demo.json文件中读取数据，并将其分成两个列表：documents和metadata。 documents是初创公司的原始文本描述。metadata是与每个初创公司相关的有效载荷，例如名称、位置和图片。 我们将使用documents将数据编码为向量。

  1. Encode and upload data.
client.add(
    collection_name="startups",
    documents=documents,
    metadata=metadata,
    parallel=0,  # Use all available CPU cores to encode data. 
    # Requires wrapping code into if __name__ == '__main__' block
)
Upload processed data

从这里下载并解压处理后的数据，或使用以下脚本：

wget https://storage.googleapis.com/dataset-startup-search/startup-list-com/startups_hybrid_search_processed_40k.tar.gz
tar -xvf startups_hybrid_search_processed_40k.tar.gz

然后你可以将数据上传到Qdrant。

from typing import List
import json
import numpy as np
from qdrant_client import models


def named_vectors(vectors: List[float], sparse_vectors: List[models.SparseVector]) -> dict:
    # make sure to use the same client object as previously
    # or `set_model_name` and `set_sparse_model_name` manually
    dense_vector_name = client.get_vector_field_name()
    sparse_vector_name = client.get_sparse_vector_field_name()  
    for vector, sparse_vector in zip(vectors, sparse_vectors):
        yield {
            dense_vector_name: vector,
            sparse_vector_name: models.SparseVector(**sparse_vector),
        } 

with open("dense_vectors.npy", "rb") as f:
    vectors = np.load(f)
    
with open("sparse_vectors.json", "r") as f:
    sparse_vectors = json.load(f)
    
with open("payload.json", "r",) as f:
    payload = json.load(f)

client.upload_collection(
    "startups", vectors=named_vectors(vectors, sparse_vectors), payload=payload
)

add 方法将编码所有文档并将它们上传到 Qdrant。 这是两个 fastembed 特定方法之一，它将编码和上传结合到一个步骤中。

parallel 参数启用数据并行而不是内置的 ONNX 并行。

此外，如果您希望以后使用它们来更新或删除文档，可以为每个文档指定ID。 如果您不指定ID，它们将自动生成并作为add方法的结果返回。

您可以通过将tqdm进度条传递给add方法来监控编码的进度。

from tqdm import tqdm

client.add(
    collection_name="startups",
    documents=documents,
    metadata=metadata,
    ids=tqdm(range(len(documents))),
)

构建搜索API

现在所有的准备工作都已完成，让我们开始构建一个神经搜索类。

为了处理传入的请求，混合搜索类需要三样东西：1）将查询转换为向量的模型，2）执行搜索查询的Qdrant客户端，3）重新排序密集和稀疏搜索结果的融合函数。

Fastembed 集成将查询编码、搜索和融合封装到一个单一的方法调用中。 Fastembed 利用倒数秩融合来结合结果。

  1. Create a file named hybrid_searcher.py and specify the following.
from qdrant_client import QdrantClient


class HybridSearcher:
    DENSE_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
    SPARSE_MODEL = "prithivida/Splade_PP_en_v1"
    def __init__(self, collection_name):
        self.collection_name = collection_name
        # initialize Qdrant client
        self.qdrant_client = QdrantClient("http://localhost:6333")
        self.qdrant_client.set_model(self.DENSE_MODEL)
        # comment this line to use dense vectors only
        self.qdrant_client.set_sparse_model(self.SPARSE_MODEL)
  1. Write the search function.
def search(self, text: str):
    search_result = self.qdrant_client.query(
        collection_name=self.collection_name,
        query_text=text,
        query_filter=None,  # If you don't want any filters for now
        limit=5,  # 5 the closest results
    )
    # `search_result` contains found vector ids with similarity scores 
    # along with the stored payload
    
    # Select and return metadata
    metadata = [hit.metadata for hit in search_result]
    return metadata
  1. Add search filters.

使用Qdrant，还可以在搜索中添加一些条件。 例如，如果你想搜索某个城市的初创公司，搜索查询可能如下所示：

from qdrant_client import models

    ...

    city_of_interest = "Berlin"

    # Define a filter for cities
    city_filter = models.Filter(
        must=[
            models.FieldCondition(
                key="city", 
                match=models.MatchValue(value=city_of_interest)
            )
        ]
    )

    search_result = self.qdrant_client.query(
        collection_name=self.collection_name,
        query_text=text,
        query_filter=city_filter,
        limit=5
    )
    ...

你现在已经为神经搜索查询创建了一个类。现在将其封装成一个服务。

使用FastAPI部署搜索

要构建服务，您将使用FastAPI框架。

  1. Install FastAPI.

要安装它，请使用命令

pip install fastapi uvicorn
  1. Implement the service.

创建一个名为 service.py 的文件并指定以下内容。

该服务将只有一个API端点，并且看起来像这样：

from fastapi import FastAPI

# The file where HybridSearcher is stored
from hybrid_searcher import HybridSearcher

app = FastAPI()

# Create a neural searcher instance
hybrid_searcher = HybridSearcher(collection_name="startups")


@app.get("/api/search")
def search_startup(q: str):
    return {"result": hybrid_searcher.search(text=q)}


if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host="0.0.0.0", port=8000)
  1. Run the service.
python service.py
  1. Open your browser at http://localhost:8000/docs.

您应该能够看到服务的调试界面。

FastAPI Swagger interface

请随意尝试，对我们的语料库中的公司进行查询，并查看结果。

加入我们的Discord社区，在这里我们讨论向量搜索和相似性学习，发布其他神经网络和神经搜索应用的示例。

这个页面有用吗？

感谢您的反馈！🙏

我们很抱歉听到这个消息。😔 你可以在GitHub上编辑这个页面，或者创建一个GitHub问题。

本页内容：