使用Sentence Transformers和Qdrant构建神经搜索服务

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

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

神经搜索服务利用人工神经网络来提高搜索结果的准确性和相关性。除了提供简单的关键词结果外，该系统还可以通过意义检索结果。它能够理解和解释复杂的搜索查询，并提供更具上下文相关性的输出，从而有效提升用户的搜索体验。

工作流程

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

注意: 本教程的代码可以在这里找到: | 第一步: 数据准备过程 | 第二步: 神经搜索的完整代码. |

先决条件

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

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

准备样本数据集

要对创业公司描述进行神经搜索，首先必须将描述数据编码为向量。要处理文本，可以使用预训练模型，如BERT或句子转换器。句子转换器库允许您方便地下载和使用许多预训练模型，例如DistilBERT、MPNet等。

First you need to download the dataset.

wget https://storage.googleapis.com/generall-shared-data/startups_demo.json

Install the SentenceTransformer library as well as other relevant packages.

pip install sentence-transformers numpy pandas tqdm

Import the required modules.

from sentence_transformers import SentenceTransformer
import numpy as np
import json
import pandas as pd
from tqdm.notebook import tqdm

你将使用一个名为all-MiniLM-L6-v2的预训练模型。这是一个性能优化的句子嵌入模型，你可以在这里了解更多关于它和其他可用模型的信息。

Download and create a pre-trained sentence encoder.

model = SentenceTransformer(
    "all-MiniLM-L6-v2", device="cuda"
)  # or device="cpu" if you don't have a GPU

Read the raw data file.

df = pd.read_json("./startups_demo.json", lines=True)

Encode all startup descriptions to create an embedding vector for each. Internally, the encode function will split the input into batches, which will significantly speed up the process.

vectors = model.encode(
    [row.alt + ". " + row.description for row in df.itertuples()],
    show_progress_bar=True,
)

所有的描述现在都已转换为向量。共有40474个384维的向量。模型的输出层具有这个维度

vectors.shape
# > (40474, 384)

Download the saved vectors into a new file named startup_vectors.npy

np.save("startup_vectors.npy", vectors, allow_pickle=False)

在 Docker 中运行 Qdrant

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

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

Download the Qdrant image from DockerHub.

docker pull qdrant/qdrant

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

Install the official Python client to best interact with Qdrant.

pip install qdrant-client

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

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

Create a client object for Qdrant.

# Import client library
from qdrant_client import QdrantClient
from qdrant_client.models import VectorParams, Distance

client = QdrantClient("http://localhost:6333")

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=VectorParams(size=384, distance=Distance.COSINE),
    )

Create an iterator over the startup data and vectors.

Qdrant 客户端库定义了一个特殊函数，允许你将数据集加载到服务中。然而，由于数据量可能太大，无法一次性装入单台计算机的内存，因此该函数接受一个数据迭代器作为输入。

fd = open("./startups_demo.json")

# payload is now an iterator over startup data
payload = map(json.loads, fd)

# Load all vectors into memory, numpy array works as iterable for itself.
# Other option would be to use Mmap, if you don't want to load all data into RAM
vectors = np.load("./startup_vectors.npy")

Upload the data

client.upload_collection(
    collection_name="startups",
    vectors=vectors,
    payload=payload,
    ids=None,  # Vector ids will be assigned automatically
    batch_size=256,  # How many vectors will be uploaded in a single request?
)

向量现已上传到Qdrant。

构建搜索API

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

为了处理传入的请求，神经搜索需要两样东西：1）一个将查询转换为向量的模型，2）用于执行搜索查询的Qdrant客户端。

Create a file named neural_searcher.py and specify the following.

from qdrant_client import QdrantClient
from sentence_transformers import SentenceTransformer


class NeuralSearcher:
    def __init__(self, collection_name):
        self.collection_name = collection_name
        # Initialize encoder model
        self.model = SentenceTransformer("all-MiniLM-L6-v2", device="cpu")
        # initialize Qdrant client
        self.qdrant_client = QdrantClient("http://localhost:6333")

Write the search function.

def search(self, text: str):
    # Convert text query into vector
    vector = self.model.encode(text).tolist()

    # Use `vector` for search for closest vectors in the collection
    search_result = self.qdrant_client.query_points(
        collection_name=self.collection_name,
        query=vector,
        query_filter=None,  # If you don't want any filters for now
        limit=5,  # 5 the most closest results is enough
    ).points
    # `search_result` contains found vector ids with similarity scores along with the stored payload
    # In this function you are interested in payload only
    payloads = [hit.payload for hit in search_result]
    return payloads

Add search filters.

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

from qdrant_client.models import Filter

    ...

    city_of_interest = "Berlin"

    # Define a filter for cities
    city_filter = Filter(**{
        "must": [{
            "key": "city", # Store city information in a field of the same name 
            "match": { # This condition checks if payload field has the requested value
                "value": city_of_interest
            }
        }]
    })

    search_result = self.qdrant_client.query_points(
        collection_name=self.collection_name,
        query=vector,
        query_filter=city_filter,
        limit=5
    ).points
    ...

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

使用FastAPI部署搜索

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

Install FastAPI.

要安装它，请使用命令

pip install fastapi uvicorn

Implement the service.

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

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

from fastapi import FastAPI

# The file where NeuralSearcher is stored
from neural_searcher import NeuralSearcher

app = FastAPI()

# Create a neural searcher instance
neural_searcher = NeuralSearcher(collection_name="startups")


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


if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host="0.0.0.0", port=8000)

Run the service.

python service.py

Open your browser at http://localhost:8000/docs.

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

FastAPI Swagger interface

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

下一步

本教程中的代码已被用于开发一个实时在线演示。您可以尝试它，以了解神经搜索在哪些情况下有用。该演示包含一个开关，可以在神经搜索和全文搜索之间进行选择。您可以打开和关闭神经搜索，以将您的结果与常规的全文搜索进行比较。