人気検索Python SDK使用例
このノートブックでは、地下鉄検索Python SDK使用方法を説明します。これは、地下鉄検索を操作するための主要なAPIとしてVectorSearchClientを提供します。
あるいは、REST APIを直接呼び出すこともできます。
要件
このノートブックは、 databricks-gte-large-enという名前のモデルサービングエンドポイントが存在することを前提としています。そのエンドポイントを作成するには、ノートブック「Mosaic AI Model Serving を使用して GTE 埋め込みモデルを呼び出す」を参照してください。
Python
%pip install --upgrade --force-reinstall databricks-vectorsearch langchain
dbutils.library.restartPython()
Python
from databricks.vector_search.client import VectorSearchClient
vsc = VectorSearchClient()
Python
help(VectorSearchClient)
おもちゃのデータセットをDeltaテーブルにロードします
以下はDeltaテーブルを作成します。
Python
# Specify the catalog and schema to use. You must have USE_CATALOG privilege on the catalog and USE_SCHEMA and CREATE_TABLE privileges on the schema.
# Change the catalog and schema here if necessary.
catalog_name = "main"
schema_name = "default"
Python
source_table_name = "en_wiki"
source_table_fullname = f"{catalog_name}.{schema_name}.{source_table_name}"
Python
# Uncomment if you want to start from scratch.
# spark.sql(f"DROP TABLE {source_table_fullname}")
Python
source_df = spark.read.parquet("/databricks-datasets/wikipedia-datasets/data-001/en_wikipedia/articles-only-parquet").limit(10)
display(source_df)
Python
source_df.write.format("delta").option("delta.enableChangeDataFeed", "true").saveAsTable(source_table_fullname)
Python
display(spark.sql(f"SELECT * FROM {source_table_fullname}"))
トレンド検索エンドポイントの作成
Python
vector_search_endpoint_name = "vector-search-demo-endpoint"
Python
vsc.create_endpoint(
name=vector_search_endpoint_name,
endpoint_type="STANDARD" # or "STORAGE_OPTIMIZED"
)
Python
endpoint = vsc.get_endpoint(
name=vector_search_endpoint_name)
endpoint
ベクターインデックスを作成する
Python
# Vector index
vs_index = "en_wiki_index"
vs_index_fullname = f"{catalog_name}.{schema_name}.{vs_index}"
embedding_model_endpoint = "databricks-gte-large-en"
Python
index = vsc.create_delta_sync_index(
endpoint_name=vector_search_endpoint_name,
source_table_name=source_table_fullname,
index_name=vs_index_fullname,
pipeline_type='TRIGGERED',
primary_key="id",
embedding_source_column="text",
embedding_model_endpoint_name=embedding_model_endpoint
)
index.describe()
ベクトルインデックスを取得する
ベクトルインデックス名を使用してベクトルインデックスオブジェクトを取得するには、 get_index()を使用します。インデックスオブジェクトに対してdescribe()使用すると、インデックスの構成情報の概要を表示することもできます。
Python
index = vsc.get_index(endpoint_name=vector_search_endpoint_name, index_name=vs_index_fullname)
index.describe()
Python
# Wait for index to come online. Expect this command to take several minutes.
import time
while not index.describe().get('status').get('detailed_state').startswith('ONLINE'):
print("Waiting for index to be ONLINE...")
time.sleep(5)
print("Index is ONLINE")
index.describe()
類似性検索
ベクターインデックスを照会して、類似の文書を検索してください。
Python
# Returns [col1, col2, ...]
# You can set this to any subset of the columns.
all_columns = spark.table(source_table_fullname).columns
results = index.similarity_search(
query_text="Greek myths",
columns=all_columns,
num_results=2)
results
Python
# Search with a filter. Note that the syntax depends on the endpoint type.
# Standard endpoint syntax
results = index.similarity_search(
query_text="Greek myths",
columns=all_columns,
filters={"id NOT": ("13770", "88231")},
num_results=2)
# Storage-optimized endpoint syntax
# results = index.similarity_search(
# query_text="Greek myths",
# columns=all_columns,
# filters='id NOT IN ("13770", "88231")',
# num_results=2)
results
結果をLangChainドキュメントに変換する
最初に取得した列はpage_contentに読み込まれ、残りはメタデータに読み込まれます。
Python
from langchain_core.documents import Document
from typing import List
def convert_vector_search_to_documents(results) -> List[Document]:
column_names = []
for column in results["manifest"]["columns"]:
column_names.append(column)
langchain_docs = []
for item in results["result"]["data_array"]:
metadata = {}
score = item[-1]
# print(score)
i = 1
for field in item[1:-1]:
# print(field + "--")
metadata[column_names[i]["name"]] = field
i = i + 1
doc = Document(page_content=item[0], metadata=metadata) # , 9)
langchain_docs.append(doc)
return langchain_docs
langchain_docs = convert_vector_search_to_documents(results)
langchain_docs
ベクターインデックスを削除
Python
vsc.delete_index(index_name=vs_index_fullname)