轻松实现基于Databricks的向量搜索:从搭建到查询的完整指南

158 阅读3分钟
# 引言

在现代数据驱动的世界中,能够快速高效地搜索和处理大规模数据变得至关重要。Databricks Vector Search是一种无服务器相似性搜索引擎,能够存储数据的向量表示及其元数据。通过与Delta表和Unity Catalog的集成,它让我们可以轻松创建并更新向量搜索索引,以便快速检索相似数据。本篇文章将带领您通过一个完整的演示来了解如何使用SelfQueryRetriever与Databricks Vector Search。

# 主要内容

## 创建Databricks向量存储索引

首先,我们需要创建一个Databricks向量存储索引并插入数据。为了这个演示,我们准备了一组有关电影摘要的小型数据集。

**注意**:使用self-query检索器,需要安装`lark`模块(使用命令`pip install lark`)。

```shell
%pip install --upgrade --quiet langchain-core databricks-vectorsearch langchain-openai tiktoken

确保重新启动内核以使用更新的软件包。

配置OpenAI Embeddings

为了使用OpenAI Embeddings,我们需要获得OpenAI API密钥。

import getpass
import os

os.environ["OPENAI_API_KEY"] = getpass.getpass("OpenAI API Key:")
databricks_host = getpass.getpass("Databricks host:")
databricks_token = getpass.getpass("Databricks token:")

初始化Vector Search客户端

接下来,我们使用Databricks Vector Search客户机来创建和管理向量搜索索引。

from databricks.vector_search.client import VectorSearchClient
from langchain_openai import OpenAIEmbeddings

embeddings = OpenAIEmbeddings()
emb_dim = len(embeddings.embed_query("hello"))

vsc = VectorSearchClient(
    workspace_url=databricks_host, personal_access_token=databricks_token
)  # 使用API代理服务提高访问稳定性
vsc.create_endpoint(name="vector_search_demo_endpoint", endpoint_type="STANDARD")

插入数据

使用事先准备的电影数据集,我们将其插入到向量搜索索引中。

from langchain_core.documents import Document

docs = [
    Document(page_content="A bunch of scientists bring back dinosaurs and mayhem breaks loose", metadata={"id": 1, "year": 1993, "rating": 7.7, "genre": "action"}),
    # 其余文档省略...
]

vector_store.add_documents(docs)

创建Self-Query Retriever

现在,我们可以实例化我们的self-query检索器。需要提前提供关于文档支持的元数据字段信息和文档内容的简短描述。

from langchain.chains.query_constructor.base import AttributeInfo
from langchain.retrievers.self_query.base import SelfQueryRetriever
from langchain_openai import OpenAI

metadata_field_info = [
    AttributeInfo(name="genre", description="The genre of the movie", type="string"),
    # 其余字段省略...
]
retriever = SelfQueryRetriever.from_llm(
    OpenAI(temperature=0), vector_store, "Brief summary of a movie", metadata_field_info, verbose=True
)

代码示例

以下代码展示了如何使用SelfQueryRetriever进行查询:

retriever.invoke("What are some movies about dinosaurs")

常见问题和解决方案

  • 网络访问限制问题:如果您在某些地区可能遇到网络访问的问题,考虑使用API代理服务以提高访问稳定性。
  • 数据更新问题:请确保在插入新数据后,索引获得更新,特别是在大规模数据环境下。

总结和进一步学习资源

本文介绍了如何使用Databricks Vector Search进行相似性搜索,从创建索引到进行复杂查询的全过程。可以进一步参考以下资源来深入学习:

参考资料

  1. Langchain Core GitHub
  2. Databricks Vector Search GitHub

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


---END---