使用Aerospike Vector Search和LangChain构建高效的语义搜索系统

标题: 使用Aerospike Vector Search和LangChain构建高效的语义搜索系统

内容:

使用Aerospike Vector Search和LangChain构建高效的语义搜索系统

在当今的AI时代,向量搜索已成为许多应用的关键功能。本文将介绍如何使用Aerospike Vector Search(AVS)和LangChain来构建一个高效的语义搜索系统。我们将通过一个实际的例子,展示如何索引和搜索大量文本数据。

1. 引言

Aerospike Vector Search(AVS)是Aerospike数据库的一个扩展,它能够在存储在Aerospike中的大型数据集上执行搜索。LangChain则是一个强大的框架,用于开发由语言模型驱动的应用程序。将这两者结合,我们可以构建出高效、可扩展的语义搜索系统。

本文将通过一个引用语句(quotes)搜索的实例,演示如何:

1. 设置AVS环境
2. 准备和加载数据
3. 创建嵌入向量
4. 构建向量索引
5. 执行相似性搜索和最大边际相关性(MMR)搜索

2. 环境设置

首先,我们需要安装必要的依赖:

pip install --upgrade --quiet aerospike-vector-search==0.6.1 langchain-community sentence-transformers langchain

注意: 由于sentence-transformers依赖较大,安装可能需要几分钟时间。

接下来,我们需要设置AVS实例的IP地址和端口:

PROXIMUS_HOST = "<avs-ip>"
PROXIMUS_PORT = 5000

3. 数据准备和加载

为了演示,我们将使用一个包含约100,000条引用语句的数据集。我们可以使用以下命令下载数据集:

wget https://github.com/aerospike/aerospike-vector-search-examples/raw/7dfab0fccca0852a511c6803aba46578729694b5/quote-semantic-search/container-volumes/quote-search/data/quotes.csv.tgz

解压后,我们使用LangChain的CSVLoader来加载数据:

import itertools
import os
import tarfile
from langchain_community.document_loaders.csv_loader import CSVLoader

filename = "./quotes.csv"

if not os.path.exists(filename) and os.path.exists(filename + ".tgz"):
    with tarfile.open(filename + ".tgz", "r:gz") as tar:
        tar.extractall(path=os.path.dirname(filename))

NUM_QUOTES = 5000
documents = CSVLoader(filename, metadata_columns=["author", "category"]).lazy_load()
documents = list(itertools.islice(documents, NUM_QUOTES))

4. 创建嵌入向量

我们使用HuggingFace的"all-MiniLM-L6-v2"模型来创建文本嵌入:

from aerospike_vector_search.types import VectorDistanceMetric
from langchain_community.embeddings import HuggingFaceEmbeddings

MODEL_DIM = 384
MODEL_DISTANCE_CALC = VectorDistanceMetric.COSINE
embedder = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")

5. 创建Aerospike索引和嵌入文档

在添加文档之前,我们需要在Aerospike数据库中创建一个索引:

from aerospike_vector_search import AdminClient, Client, HostPort
from langchain_community.vectorstores import Aerospike

seed = HostPort(host=PROXIMUS_HOST, port=PROXIMUS_PORT)
NAMESPACE = "test"
INDEX_NAME = "quote-miniLM-L6-v2"
VECTOR_KEY = "vector"

client = Client(seeds=seed)
admin_client = AdminClient(seeds=seed)

# 检查索引是否存在,如果不存在则创建
index_exists = False
for index in admin_client.index_list():
    if index["id"]["namespace"] == NAMESPACE and index["id"]["name"] == INDEX_NAME:
        index_exists = True
        print(f"{INDEX_NAME} already exists. Skipping creation")
        break

if not index_exists:
    print(f"{INDEX_NAME} does not exist. Creating index")
    admin_client.index_create(
        namespace=NAMESPACE,
        name=INDEX_NAME,
        vector_field=VECTOR_KEY,
        vector_distance_metric=MODEL_DISTANCE_CALC,
        dimensions=MODEL_DIM,
        index_meta_data={
            "model": "miniLM-L6-v2",
            "date": "05/04/2024",
            "dim": str(MODEL_DIM),
            "distance": "cosine",
        },
    )

admin_client.close()

# 将文档添加到索引中
docstore = Aerospike.from_documents(
    documents,
    embedder,
    client=client,
    namespace=NAMESPACE,
    vector_key=VECTOR_KEY,
    index_name=INDEX_NAME,
    distance_strategy=MODEL_DISTANCE_CALC,
)

6. 执行搜索

现在我们可以执行各种类型的搜索了。

6.1 相似性搜索

query = "A quote about the beauty of the cosmos"
docs = docstore.similarity_search(
    query, k=5, index_name=INDEX_NAME, metadata_keys=["_id", "author"]
)

def print_documents(docs):
    for i, doc in enumerate(docs):
        print("~~~~ Document", i, "~~~~")
        print("auto-generated id:", doc.metadata["_id"])
        print("author: ", doc.metadata["author"])
        print(doc.page_content)
        print("~~~~~~~~~~~~~~~~~~~~\n")

print_documents(docs)

6.2 最大边际相关性(MMR)搜索

MMR搜索可以帮助我们找到与查询相似但彼此不同的结果:

query = "A quote about our favorite four-legged pets"
retriever = docstore.as_retriever(
    search_type="mmr", search_kwargs={"fetch_k": 20, "lambda_mult": 0.7}
)
matched_docs = retriever.invoke(query)

print_documents(matched_docs)

6.3 带相关性阈值的搜索

我们可以设置一个相关性阈值,只返回超过该阈值的结果:

query = "A quote about stormy weather"
retriever = docstore.as_retriever(
    search_type="similarity_score_threshold",
    search_kwargs={"score_threshold": 0.4},
)
matched_docs = retriever.invoke(query)

print_documents(matched_docs)

7. 清理资源

完成操作后,请确保关闭客户端以释放资源:

client.close()

8. 总结和进一步学习

通过本文,我们展示了如何使用Aerospike Vector Search和LangChain构建一个高效的语义搜索系统。这个系统可以处理大量文本数据,并支持多种搜索方式,包括相似性搜索、MMR搜索和带阈值的搜索。

对于那些希望深入了解这个主题的读者,以下是一些进一步学习的资源:

1. Aerospike Vector Search官方文档
2. LangChain文档中的向量存储指南
3. Sentence Transformers库文档
4. 向量搜索和语义相似度的学术论文

参考资料

1. Aerospike Vector Search Documentation: https://docs.aerospike.com/vector-search
2. LangChain Documentation: https://python.langchain.com/docs/modules/data_connection/vectorstores/
3. Sentence Transformers: https://www.sbert.net/
4. “Efficient and robust approximate nearest neighbor search using Hierarchical Navigable Small World graphs” by Yu. A. Malkov, D. A. Yashunin

如果这篇文章对你有帮助,欢迎点赞并关注我的博客。您的支持是我持续创作的动力！

—END—