探索Databricks向量搜索：构建强大的相似性检索引擎

用户331287887593
61 阅读3分钟

引言

在当今的大数据时代，能够快速有效地搜索和推荐相似的数据是许多应用的核心需求。Databricks Vector Search是一种无服务器的相似性搜索引擎，它允许您将数据的向量表示（包括元数据）存储在向量数据库中。利用Vector Search，您可以从由Unity Catalog管理的Delta表中创建自动更新的向量搜索索引，并使用简单的API查询返回最相似的向量。在本文中，我们将演示如何使用Databricks Vector Search构建一个简单而强大的电影相似性搜索引擎。

主要内容

创建Databricks向量存储索引

首先，我们需要创建一个Databricks向量存储索引并插入一些数据。在这个演示中，我们将使用一组包含电影摘要的小型数据集。为了使用OpenAI的嵌入，我们需要获取OpenAI API密钥。

安装必要的库

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

注意: 安装完库后可能需要重启内核以加载更新的包。

设置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:")

配置向量搜索客户端

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

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

vector_search_endpoint_name = "vector_search_demo_endpoint"

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

创建索引并描述

index_name = "demo_index"

index = vsc.create_direct_access_index(
    endpoint_name=vector_search_endpoint_name,
    index_name=index_name,
    primary_key="id",
    embedding_dimension=emb_dim,
    embedding_vector_column="text_vector",
    schema={
        "id": "string",
        "page_content": "string",
        "year": "int",
        "rating": "float",
        "genre": "string",
        "text_vector": "array<float>",
    },
)

index.describe()

代码示例

接下来，我们演示如何使用创建的索引进行查询。

from langchain_core.documents import Document
from langchain_community.vectorstores import DatabricksVectorSearch

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 = DatabricksVectorSearch(
    index,
    text_column="page_content",
    embedding=embeddings,
    columns=["year", "rating", "genre"],
)

vector_store.add_documents(docs)

# 创建自查询检索器
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"),
    AttributeInfo(name="year", description="The year the movie was released", type="integer"),
    AttributeInfo(name="rating", description="A 1-10 rating for the movie", type="float")
]
document_content_description = "Brief summary of a movie"
llm = OpenAI(temperature=0)

retriever = SelfQueryRetriever.from_llm(
    llm, vector_store, document_content_description, metadata_field_info, verbose=True
)

# 测试检索器
retriever.invoke("What are some movies about dinosaurs")

常见问题和解决方案

  • API访问受限：由于某些地区的网络限制，开发者可以考虑使用API代理服务来提高访问的稳定性。
  • 数据量大导致查询效率低：优化索引结构，并合理使用元数据过滤可以有效提升查询效率。

总结和进一步学习资源

Databricks Vector Search为开发者提供了一种强大的工具，可以方便地建立高效的相似性检索系统。结合OpenAI的嵌入模型，您可以实现多种复杂的查询和过滤逻辑。进一步的学习可以访问以下资源：

参考资料

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

---END---

avatar
文章
阅读
粉丝