社区首页 >专栏 >构建自然语言搜索引擎

构建自然语言搜索引擎

杜逸先

发布于 2023-12-25 02:03:14

44500

代码可运行

文章被收录于专栏：追不上乌龟的兔子追不上乌龟的兔子

运行总次数：0

代码可运行

近期简单学习了一下向量数据库 qdrant 与 sentence-transformers 库，两者结合可以构建一个简单的自然语言搜索引擎。顺着官方的教程实操了一遍之后，稍微调整一番，我在中文数据集上构建了一个自然语言搜索引擎。

数据采集

教程中的数据集是一些美国的初创公司的数据（来自 startups-list.com），我打算自己从互联网上采集一些中文数据集，正好近期一直在使用微信读书，决定采集微信读书平台上计算机分类下的书籍数据。

观察了微信读书网页版的请求之后，我发现可以通过 ‘https://weread.qq.com/web/bookListInCategory/{type}?maxIndex={maxIndex}‘ 获取指定分类下的书籍列表。在实验过程中，我发现即使分类下有更多的书籍，maxIndex 超过 480 之后就不再返回新的数据了，这应该是微信读书官方的一些限制。不过计算机分类下一共有七个子分类，每个分类都采集四百多本书籍，总共有 3000 多本书籍，也足够使用了。

首先定义一个函数用于获取制定分类指定页数的书籍列表。

url = 'https://weread.qq.com/web/bookListInCategory/'
def fetch_page(param):
    type, page = param
    _url = f"{url}{type}?maxIndex={page*20}"
    resp = requests.get(_url)
    return resp.json()["books"]

为了提高采集的效率，我创建了一个线程池用于批量执行采集任务。

from itertools import product
with ThreadPoolExecutor(max_workers=10) as executor:
    pages = list(
        executor.map(
            fetch_page,
            product(
                [700001, 700002, 700003, 700004, 700005, 700006, 700007],
                range(1, 25),
            ),
        )
    )

books = []
for page in pages:
    books.extend(page)

最终的执行时间是 11 秒，相当不错的速度。

数据处理

参考教程里的流程，我使用 pandas 进行了初步的数据处理，并用 sentence-transformers 库生成书籍描述的向量表示。

from sentence_transformers import SentenceTransformer
import numpy as np
import pandas as pd

model = SentenceTransformer(
    "paraphrase-multilingual-MiniLM-L12-v2", device="cuda"
)

df = pd.DataFrame([book['bookInfo'] for book in books])

vectors = model.encode(
    [row.title + ". " + row.author + "." + row.intro for row in df.itertuples()],
    show_progress_bar=True,
)

vectors.shape  # (3244, 384)

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

使用 qdrant 持久化向量数据

终于到了 qdrant 的环节，这里简单介绍一下 qdrant（内容来自 Microsoft Copilot）。

Qdrant 是一个开源的向量数据库和向量相似度搜索引擎，用 Rust 语言编写，可以快速、可靠地存储和搜索任意维度的向量，支持多种距离度量，如余弦、欧氏、曼哈顿等。Qdrant 可以用于构建基于语义嵌入或神经网络编码器的匹配、搜索、推荐等 AI 应用。

Qdrant 的主要特点有：

- 易于使用的 API，提供 OpenAPI v3 规范，可以生成几乎任何编程语言的客户端库，或者使用现成的 Python 或其他语言的客户端。
- 快速和准确，采用一种独特的自定义修改的 HNSW 算法进行近似最近邻搜索，具有最先进的速度，并且可以在不影响结果的情况下应用搜索过滤器。
- 可过滤，支持与向量关联的额外负载，不仅存储负载，而且还允许根据负载值过滤结果。与 Elasticsearch 的后过滤不同，Qdrant 保证检索到所有相关的向量。
- 丰富的数据类型，向量负载支持多种数据类型和查询条件，包括字符串匹配、数值范围、地理位置等。负载过滤条件允许你构建几乎任何应该在相似度匹配之上工作的自定义业务逻辑。
- 分布式，支持云原生和水平扩展。无论你需要服务多少数据，Qdrant 都可以使用合适的计算资源。
- 高效，有效地利用你的资源。完全用 Rust 语言开发，实现了动态查询规划和负载数据索引。硬件感知的构建也可用于企业。

参照官方文档的指导，我使用 docker 部署了一个 qdrant 服务。

docker run -d -p 6333:6333 -v $(pwd)/qdrant_data:/data qdrant/qdrant

客户端使用的是官方提供的 python sdk。

pip install qdrant-client

创建集合

连接到 qdrant 服务之后，创建一个集合用于存储书籍数据。

from qdrant_client import QdrantClient
from qdrant_client.models import VectorParams, Distance

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

qdrant_client.recreate_collection(
    collection_name="wx_books",
    vectors_config=VectorParams(size=384, distance=Distance.COSINE),
)

这里指定了向量的维度为 384，距离度量方法为余弦相似度。

导入数据

接下来使用 upload_collection 方法，将书籍详情与向量数据导入到 qdrant 中。

# vectors = np.load("./wx_books.npy")
# books = json.load(open("./wx_books.json", "r", encoding="utf-8"))

qdrant_client.upload_collection(
    collection_name="wx_books",
    vectors=vectors,
    payload=(book['bookInfo'] for book in books),
    ids=None,  # Vector ids will be assigned automatically
    batch_size=256,  # How many vectors will be uploaded in a single request?
)

这三千多条书籍数据的导入在我本地花费了 2 秒钟左右，速度还是很快的。

构建搜索引擎

这里可以完全照搬官方教程，创建一个 NeuralSearcher 类，用于在 qdrant 集合中进行自然语言搜索。

# search.py

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(
            "paraphrase-multilingual-MiniLM-L12-v2", device="cpu"
        )
        # initialize Qdrant client
        self.qdrant_client = QdrantClient("http://localhost:6333")

    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.search(
            collection_name=self.collection_name,
            query_vector=vector,
            query_filter=None,  # If you don't want any filters for now
            limit=5  # 5 the most closest results is enough
        )
        # `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

接下来使用 FastAPI 创建一个简单的接口，用于接收用户的查询请求。

# api.py
from fastapi import FastAPI

# The file where NeuralSearcher is stored
from search import NeuralSearcher

app = FastAPI()

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


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

效果

定义一个 search 函数用来测试搜索服务的效果。

def search(q: str):
    resp = requests.get(f"http://localhost:8000/api/search?q={q}")
    for book in resp.json()['result']:
        print(f'{book["title"]} - {book["author"]}')
        print(book['intro'][:50])
        print()

既然主打的是自然语言搜索，查询条件可以不用太简单（例如“AI”，“C++”等关键词。）

search("初学者如何入门机器学习")

机器学习入门：数学原理解析及算法实践 - 董政
本书面向初学者，介绍了机器学习的基本方法，循序渐进的阐述了其中的数学原理，让读者能够知其然，然后知其

机器学习：Python实践 - 魏贞原
本书系统地讲解了机器学习的基本知识，以及在实际项目中使用机器学习的基本步骤和方法；详细地介绍了在进行

机器学习实践指南：案例应用解析 - 麦好
《机器学习实践指南：案例应用解析》是机器学习及数据分析领域不可多得的一本著作，也是为数不多的既有大量

机器学习：使用OpenCV、Python和scikit-learn进行智能图像处理（原书第2版） - 阿迪蒂亚·夏尔马 维什韦什·拉维·什里马利 迈克尔·贝耶勒
本书通过具体的编程实践案例，全面系统地讲述了机器学习涉及的核心内容。首先介绍新特性以及安装OpenC

机器学习 - 赵卫东 董亮
机器学习是人工智能的重要技术基础，涉及的内容十分广泛。本书内容涵盖了机器学习的基础知识，主要包括机器

search("我想了解生成式人工智能（AIGC）")

AIGC未来已来 - 翟尤 郭晓静 曾宣玮
AIGC（Artificial Intelligence Generated Content）中文译

人工智能基础与进阶（第二版） - 周越编著
人工智能是一门发展极其迅速且内容丰富的学科，其众多分支领域都值得大家去探索和学习。本书分为基础篇和进

人工智能（AI）应用从入门到精通 - 苏秉华 吴红辉 滕悦然编著
《人工智能（AI）应用从入门到精通》是一本人工智能应用入门级读物，全书分基础篇和应用篇两个部分。基础

人工智能简史 - 尼克
本书全面讲述人工智能的发展史，几乎覆盖人工智能学科的所有领域，包括人工智能的起源、自动定理证明、专家

生成式AI：人工智能的未来 - 詹姆斯·斯金纳
一本书全面了解生成式AI的发展与创作能力，并为我所用。20世纪60年代，AI的概念就被提出，其商业应