向量数据库入门（一）文本相似性搜索

本文将介绍三种常用的向量数据库：faiss, Milvus和Qdrant，并给出一个具体的使用例子。

向量数据库（Vector Database）是一种专门用于存储、管理、查询、检索向量的数据库，主要应用于人工智能、机器学习、数据挖掘等领域。

在向量数据库中，数据以向量的形式进行存储和处理，需要将原始的非向量型数据转化为向量表示（比如文本使用Embedding技术获得其表征向量）。

这种数据库能够高效地进行相似性搜索，快速找到最相似的向量，适用于人脸识别、图像搜索、推荐系统等需要相似性匹配的应用。

去年大模型的大火也带动了向量数据库的迅速发展，使得向量数据库成为热门方向之一，成为AI领域不可或缺的一项重要工具。

本文将介绍常见的三种向量数据库，并结合具体的样例文本给出它们在文本相似性搜索方面的应用。这三个向量数据库分别为：

faiss
Milvus
Qdrant

获取文本Embedding向量

我们使用的样例文本来源于百度百科的“中国载人登月工程”词条，将其本文切分为句子，过滤其中的纯数字的句子，保存为dengyue.txt文件。

对于上述文本，使用OpenAI的embedding模型text-embedding-ada-002来获取句子的表征向量（向量维度为1536，L2范数为1，即单位向量），并保存为numpy模块的npz文件，用于离线存储，避免每次使用时都需要加载。

示例的代码如下：

# -*- coding: utf-8 -*-
import os
import re
import json
from typing import List

from dotenv import load_dotenv
import requests
import numpy as np
from sentencex import segment

load_dotenv()


def get_embedding(texts: List[str]):
    url = "https://api.openai.com/v1/embeddings"
    payload = json.dumps({
        "model": "text-embedding-ada-002",
        "input": texts,
        "encoding_format": "float"
    })
    headers = {
        'Content-Type': 'application/json',
        'Authorization': f'Bearer {os.getenv("OPENAI_API_KEY")}'
    }
    response = requests.request("POST", url, headers=headers, data=payload)
    embedding = [_["embedding"] for _ in response.json()['data']]
    response.close()
    return embedding


def read_file():
    with open("dengyue.txt", "r", encoding="utf-8") as f:
        content = [_.strip() for _ in f.readlines() if _.strip()]
    file_sentences = []
    for line in content:
        sents = list(segment(language="zh", text=line))
        for sent in sents:  # filter
            if sent and not re.match(r'\[\d+\]', sent):
                file_sentences.append(sent)
    return file_sentences


if __name__ == '__main__':
    sentences = read_file()
    sentences_embeddings = np.array(get_embedding(sentences))
    np.savez('text_embedding.npz', sentences_embeddings)
    with open('text.json', 'w') as f:
        f.write(json.dumps({i: sentence for i, sentence in enumerate(sentences)}, indent=4, ensure_ascii=False))
    # sentences_embeddings = np.load('text_embedding.npz')
    # print(sentences_embeddings['arr_0'][0].tolist())

运行上述文件，则会得到句子的表征向量文件text_embedding.npz，使用np.load()方法即可加载这些向量。

faiss

faiss全称为Facebook AI Similarity Search，是Facebook开源的一个用于高效相似性搜索和密集向量聚类的库，使用C++编写，提供python接口，对10亿量级的索引可以做到毫秒级检索的性能。它包含的算法可以搜索任意大小的向量集，甚至可能无法容纳在 RAM 中的向量集。它还包含用于评估和参数调整的支持代码。

faiss模块同时支持CPU和GPU.

faiss支持的索引(Index)类型如下图：

文本相似性搜索：获取与查询文本（query）相似度最接近的k个语料文本：

获取query的表征向量（同一个Embedding模型）
借助faiss查询出与该query向量相似度最接近的k个向量
获取这k个向量对应的文本

示例Python代码如下：

# -*- coding: utf-8 -*-
import json
import numpy as np
from faiss import IndexFlatIP

from get_text_embedding import get_embedding

with open('text.json', 'r') as f:
    sentences = list(json.loads(f.read()).values())

# create faiss index
faiss_index = IndexFlatIP(1536)
# add vector
sentences_embeddings = np.load("text_embedding.npz")['arr_0']
faiss_index.add(sentences_embeddings)
# vector search
query = "阿波罗登月计划"
query = "神舟五号 杨利伟"
# query = "北京航天城"
query_embedding = get_embedding([query])
distances, doc_indices = faiss_index.search(np.array(query_embedding), 3)
# result
print(f"query: {query}")
print("recall result: ")
for i, sent_index in enumerate(doc_indices.tolist()[0]):
    print(f"rank: {i}, similarity: {distances.tolist()[0][i]}, text: {sentences[sent_index]}")
# reset index
faiss_index.reset()

输入query为阿波罗登月计划，最接近的3个语料文本为：

query: 阿波罗登月计划
recall result: 
rank: 0, similarity: 0.855633020401001, text: 二十世纪六、七十年代，美国人凭借“阿波罗”计划率先登上月球。
rank: 1, similarity: 0.8497034907341003, text: 初步方案是：采用两枚运载火箭分别将月面着陆器和载人飞船送至环月轨道在轨交会对接，航天员从飞船进入月面着陆器。
rank: 2, similarity: 0.8401131629943848, text: 欧阳自远表示，对于载人登月工程国家还没有公布一个明确的时间表，但已经在做积极的论证和准备。

输入query为神舟五号杨利伟，最接近的3个语料文本为：

query: 神舟五号 杨利伟
recall result: 
rank: 0, similarity: 0.8733900785446167, text: 叶培建表示，中国载人航天计划共有“三步走”，第一步是从“神舟一号”走到“神舟七号”，前4个载人飞船是无人的，“神五”时杨利伟上天，实现中国人第一次进入太空的梦想，虽然比“太空旅行第一人”尤里·阿列克谢耶维奇·加加林晚了几十年，但两个不是一个数量级。
rank: 1, similarity: 0.8618367910385132, text: 加加林只飞了一个小时，是伞降；杨利伟飞了一天，是软着陆，返回舱已达到“容3个人7天往返”标准。
rank: 2, similarity: 0.8449346423149109, text: 加加林落地偏差达400公里；杨利伟落地时，落在中国内蒙古，当时设计“100乘以90公里着陆”都算准确着陆，但最终落地偏差只有10公里，后来“神六”、“神七”落地误差只有1公里，3条飞船实现了第一阶段的全部任务。

输入query为北京航天城，最接近的3个语料文本为：

query: 北京航天城
recall result: 
rank: 0, similarity: 0.8499420881271362, text: 在北京航天城当天举行的“李锦记航天奖学金”颁发仪式上，戚发轫介绍，根据中国的规划，2014年左右在深空探测领域，将把十几吨的航天器送到地球轨道，2020年前将建成自己的空间站。
rank: 1, similarity: 0.8408806920051575, text: 5月29日上午，神舟十六号载人飞行任务新闻发布会在酒泉卫星发射中心举行。
rank: 2, similarity: 0.8388772010803223, text: 戚发轫还介绍，中国将建成全球的自主北斗导航系统，中国的人、车、船、飞机在世界任何地方都可靠自己的卫星来导航定位。

Milvus

Milvus是在2019年创建的，其唯一目标是存储、索引和管理由深度神经网络和其他机器学习（ML）模型生成的大规模嵌入向量。

作为一个专门设计用于处理输入向量查询的数据库，它能够处理万亿级别的向量索引。与现有的关系型数据库主要处理遵循预定义模式的结构化数据不同，Milvus从底层设计用于处理从非结构化数据转换而来的嵌入向量。

我们使用Docker启动Milvus服务，并进行文本相似性搜索。示例Python代码如下：

# -*- coding: utf-8 -*-
# reference: https://milvus.io/docs
import json
import time
import numpy as np
from pymilvus import MilvusClient, FieldSchema, CollectionSchema, DataType, Collection

from get_text_embedding import get_embedding

with open('text.json', 'r') as f:
    sentences = list(json.loads(f.read()).values())

# add vector
sentences_embeddings = np.load("text_embedding.npz")['arr_0']

collection_name = "dengyue"
# Connects to a server
client = MilvusClient(uri="http://localhost:19530", db_name="default")
# List all collection names
collections = client.list_collections()
print("exist collections: ", collections)

# create collection if not exists
if collection_name not in collections:
    # Creates a collection
    fields = [
        FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=False),
        FieldSchema(name="text", dtype=DataType.VARCHAR, max_length=1000),
        FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=1536)
    ]
    schema = CollectionSchema(fields, "text search demo")
    index_params = client.prepare_index_params()
    index_params.add_index(
        field_name="embedding",
        index_type="IVF_FLAT",
        metric_type="IP",
        params={"nlist": 128}
    )
    client.create_collection(
        collection_name=collection_name,
        schema=schema,
        index_params=index_params
    )
    time.sleep(3)
    res = client.get_load_state(
        collection_name=collection_name
    )
    print("load state: ", res)

    # Inserts vectors in the collection
    entities = [
        {"id": i + 1,
         "text": sentences[i],
         "embedding": sentences_embeddings[i].tolist()}
        for i in range(len(sentences))
    ]
    client.insert(collection_name=collection_name, data=entities)

# Single vector search
query = "阿波罗登月计划"
# query = "神舟五号 杨利伟"
# query = "北京航天城"
query_embedding = get_embedding([query])
print(query_embedding[0])
res = client.search(
    collection_name=collection_name,
    data=query_embedding,
    limit=3,
    search_params={"metric_type": "IP", "params": {}},
    output_fields=['text']
)

# Convert the output to a formatted JSON string
result = json.dumps(res, indent=4, ensure_ascii=False)
print(result)

其中创建的Milvus collection为dengyue，其schema如下（可视化工具采用Attu）：

运行上述程序，得到的文本相似度搜索结果同上述faiss。

Qdrant

Qdrant是一个用 Rust 编写的开源矢量数据库和矢量相似性搜索引擎，旨在利用先进的高性能矢量相似性搜索技术为下一代人工智能应用赋能。它的主要功能包括多语言支持（可实现各种数据类型的通用性）和适用于各种应用的过滤器。

Qdrant对近似近邻搜索的HNSW算法进行了定制修改，通过保持精确的结果，确保了最先进的搜索能力，从而实现了速度和精确度的双赢。此外，它还支持与向量相关的额外有效载荷，允许根据有效载荷值过滤结果。Qdrant 支持丰富的数据类型和查询条件，包括字符串匹配、数值范围和地理位置，为数据管理提供了多功能性。

Qdrant服务的启动方式如下：

本地模式，不需要服务器（离线db文件或内存）
本地服务器部署
Qdrant云

以本地模式（离线db文件）启动Qdrant服务，进行文本相似度搜索。示例Python代码如下：

# -*- coding: utf-8 -*-
import json
import os
import numpy as np
from qdrant_client import QdrantClient
from qdrant_client.http.models import Distance, VectorParams
from qdrant_client.http.models import PointStruct

from get_text_embedding import get_embedding


db_dir = "qdrant_db"
collect_name = "dengyue"
client = QdrantClient(path=db_dir)


if not os.path.exists(f"./{db_dir}/collection/{collect_name}"):
    with open('text.json', 'r') as f:
        sentences = list(json.loads(f.read()).values())
    sentences_embeddings = np.load('text_embedding.npz')
    # create collection
    client.recreate_collection(
        collection_name=collect_name,
        vectors_config=VectorParams(size=1536, distance=Distance.COSINE),
    )
    # insert data into collection
    operation_info = client.upsert(
        collection_name=collect_name,
        wait=True,
        points=[
            PointStruct(
                id=i + 1,
                vector=sentences_embeddings['arr_0'][i].tolist(),
                payload={
                    "text": sentences}) for i,
            sentences in enumerate(sentences)])
    print(operation_info)


# query = "阿波罗登月计划"
# query = "神舟五号 杨利伟"
query = "北京航天城"
query_embedding = get_embedding([query])[0]
# search vector
search_result = client.search(
    collection_name=collect_name,
    query_vector=query_embedding,
    limit=3
)
for search_item in search_result:
    print(search_item)

启动上述程序，会在本地同级目录下生成qdrant_db/collection/dengyue/storage.sqlite文件，表明该服务使用sqlite作为后端数据库。

运行上述程序，得到的文本相似度搜索结果同上述faiss。

Qdrant还支持本地Docker启动，启动命令为：

1	`docker run -itd -p 6333:6333 -p 6334:6334 -v $(pwd)/qdrant_storage:/qdrant/storage:z qdrant/qdrant`

此时，连接Qdrant的Python代码改为：

1	`client = QdrantClient("localhost", port=6333)`

与此同时，该服务还提供了WebUI，支持在浏览器端进行操作，并提供了RESTful风格的Console界面。

运行命令GET collections，结果如下：

{
  "result": {
    "collections": [
      {
        "name": "test_collection"
      },
      {
        "name": "dengyue"
      }
    ]
  },
  "status": "ok",
  "time": 0.000011666
}