Vector Search 快速入门

在 Gemini Enterprise Agent Platform Vector Search 快速入门中,了解如何根据虚构电子商务服装网站的示例数据集创建索引。在本快速入门中,已经创建了嵌入。本快速入门旨在帮助您在 30 分钟内开始创建和部署索引。

前提条件

本教程需要与结算账号关联的 Google Cloud 项目。如需创建新项目,请参阅设置项目和开发环境。 您需要创建项目并设置结算账号。

选择运行时环境

本教程可以在 Colab 或 Gemini Enterprise Agent Platform Workbench 上运行。

  • Colab:在 Colab 中打开本教程
  • Gemini Enterprise Agent Platform Workbench:在 Gemini Enterprise Agent Platform Workbench 中打开本教程。如果这是您第一次在 Google Cloud 项目中使用 Gemini Enterprise Agent Platform Workbench,请前往 Google Cloud 控制台的 Gemini Enterprise Agent Platform Workbench 部分,然后点击启用以启用 Notebooks API。

如要在 GitHub 中查看此笔记本,请参阅 GitHub

完成本快速入门的费用

完成本教程大约花费几美元。以下页面提供了本教程中使用的 Google Cloud 服务的价格:

您还可以使用价格计算器根据预计使用量计算估算费用。

设置

在开始使用 Gemini Enterprise Agent Platform 之前,您需要设置以下内容:

安装 Agent Platform SDK for Python

您可以通过多种方式访问 Gemini Enterprise Agent Platform 和 Cloud Storage API,包括 REST API 和 Python 版 Agent Platform SDK。本教程使用 Agent Platform SDK for Python。

!pip install --upgrade --user google-cloud-aiplatform>=1.29.0 google-cloud-storage

如需在此 Jupyter 运行时中使用新安装的软件包,您需要重启运行时,如以下代码段所示。

# Restart kernel after installs so that your environment can access the new packages
import IPython

app = IPython.Application.instance()
app.kernel.do_shutdown(True)

环境变量

设置环境变量。如果系统提示,请将 your-project-id 替换为您的项目 ID 并运行该单元。

# get project ID
PROJECT_ID = ! gcloud config get-value project
PROJECT_ID = PROJECT_ID[0]
LOCATION = "us-central1"
if PROJECT_ID == "(unset)":
    print(f"Please set the project ID manually below")
# define project information
if PROJECT_ID == "(unset)":
  PROJECT_ID = "[your-project-id]"

# generate a unique id for this session
from datetime import datetime
UID = datetime.now().strftime("%m%d%H%M")

身份验证(仅限 Colab)

如果您在 Colab 上运行此笔记本,则需要运行以下单元身份验证。如果您使用的是 Gemini Enterprise Agent Platform Workbench,则无需执行此步骤,因为它已预先进行身份验证。

import sys

# if it's Colab runtime, authenticate the user with Google Cloud
if 'google.colab' in sys.modules:
    from google.colab import auth
    auth.authenticate_user()

设置 IAM 权限

您需要向默认服务账号添加访问权限,才能使用这些服务。

  1. 前往 Google Cloud 控制台中的 IAM 页面。
  2. 查找默认计算服务账号的主账号。代码应如下所示:compute@developer.gserviceaccount.com
  3. 点击修改按钮,并为默认计算服务账号授予以下角色:Gemini Enterprise Agent Platform User、Storage Admin 和 Service Usage Admin。

启用 API

运行以下命令,为此 Google Cloud 项目启用适用于 Compute Engine、Gemini Enterprise Agent Platform 和 Cloud Storage 的 API。

! gcloud services enable compute.googleapis.com aiplatform.googleapis.com storage.googleapis.com --project {PROJECT_ID}

准备示例数据

在本教程中,我们使用 TheLook 数据集,该数据集具有一个产品表,其中包含虚构电子商务服装网站的大约 5,000 行合成产品数据。

示例数据集

我们在此表中准备了 product-embs.json 文件。

产品嵌入示例

此文件采用 JSONL 格式,每行都有 ID(商品 ID)、名称(商品名称)以及嵌入(之前使用 Gemini Enterprise Agent Platform 文本嵌入生成的 768 个维度中的商品名称嵌入)。

文本嵌入代表服饰商品名称的含义。在本教程中,我们使用 Vector Search 来完成项的语义搜索。此示例代码可用作其他快速推荐系统的基础,您可以在其中快速找到“与此产品类似的其他产品”。

如需详细了解如何根据 BigQuery 表中的数据创建嵌入,并将其存储在 JSON 文件中,请参阅文本嵌入 + Gemini Enterprise Agent Platform Vector Search 使用入门

在 Cloud Storage 上准备数据

如需使用 Gemini Enterprise Agent Platform 构建索引,请将嵌入文件放在 Cloud Storage 存储桶中。以下代码完成两个任务:

  1. 创建 Cloud Storage 存储桶。
  2. 将示例文件复制到 Cloud Storage 存储桶。
BUCKET_URI = f"gs://{PROJECT_ID}-vs-quickstart-{UID}"
! gcloud storage buckets create $BUCKET_URI --location=$LOCATION --project=$PROJECT_ID
! gcloud storage cp "gs://github-repo/data/vs-quickstart/product-embs.json" $BUCKET_URI

如需使用 Vector Search 运行查询,您还需要将嵌入文件复制到本地目录:

! gcloud storage cp "gs://github-repo/data/vs-quickstart/product-embs.json" . # for query tests

构建和部署 Vector Search 索引

了解如何创建索引,创建索引端点,然后将索引部署到该端点。

创建索引

现在可以将嵌入加载到 Vector Search 中。API 可在 SDK 的 aiplatform 软件包下获得。

# init the aiplatform package
from google.cloud import aiplatform
aiplatform.init(project=PROJECT_ID, location=LOCATION)

使用 create_tree_ah_index 函数创建 MatchingEngineIndex(Matching Engine 是 Vector Search 的旧名称)。

# create Index
my_index = aiplatform.MatchingEngineIndex.create_tree_ah_index(
    display_name = f"vs-quickstart-index-{UID}",
    contents_delta_uri = BUCKET_URI,
    dimensions = 768,
    approximate_neighbors_count = 100,
)

MatchingEngineIndex.create_tree_ah_index() 方法会构建索引。如果数据集很小,则此过程需要不到 10 分钟;否则需要约 60 分钟或更长时间,具体取决于数据集的大小。您可以在 Vector Search Google Cloud 控制台中检查索引创建状态

请参阅“索引”

用于创建索引的参数:

  • contents_delta_uri:存储嵌入 JSON 文件的 Cloud Storage 目录的 URI
  • dimensions:每个嵌入的维度大小。在这种情况下,该值为 768,因为您使用的是文本嵌入 API 中的嵌入。
  • approximate_neighbors_count:您要在典型情况下检索的类似项数

如需详细了解如何创建索引和可用参数,请参阅创建和管理索引

创建索引端点并部署索引

如要使用索引,您需要创建一个索引端点。它充当接受索引的查询请求的服务器实例。

## create `IndexEndpoint`
my_index_endpoint = aiplatform.MatchingEngineIndexEndpoint.create(
    display_name = f"vs-quickstart-index-endpoint-{UID}",
    public_endpoint_enabled = True
)

使用索引端点,通过指定唯一的已部署索引 ID 来部署索引。

DEPLOYED_INDEX_ID = f"vs_quickstart_deployed_{UID}"
# deploy the Index to the Index Endpoint
my_index_endpoint.deploy_index(
    index = my_index, deployed_index_id = DEPLOYED_INDEX_ID
)

如果这是您第一次将此索引部署到索引端点,则大约需要 30 分钟才能自动构建并启动后端。如需查看索引部署的状态,请在Google Cloud 控制台的 Agent Platform 部分中前往部署和使用部分。选择索引

请参阅“索引”

使用 Vector Search 运行查询

在以下代码中,它会查找指定产品名称的嵌入,并使用 Vector Search 查找类似的产品名称。

获取嵌入以运行查询

首先,加载嵌入 JSON 文件,以构建产品名称和嵌入的 dict

import json

# build dicts for product names and embs
product_names = {}
product_embs = {}
with open('product-embs.json') as f:
    for l in f.readlines():
        p = json.loads(l)
        id = p['id']
        product_names[id] = p['name']
        product_embs[id] = p['embedding']

使用 product_embs 字典,您可以指定产品 ID 以获取其嵌入。

 # Get the embedding for ID 6523 "cloudveil women's excursion short"
 # You can also try with other IDs such as 12711, 18090, 19536 and 11863
query_emb = product_embs['6523']

运行查询

将嵌入传递给 Endpoint.find_neighbors() 方法以查找类似的产品名称。

# run query
response = my_index_endpoint.find_neighbors(
    deployed_index_id = DEPLOYED_INDEX_ID,
    queries = [query_emb],
    num_neighbors = 10
)

# show the results
for idx, neighbor in enumerate(response[0]):
    print(f"{neighbor.distance:.2f} {product_names[neighbor.id]}")

借助 ScaNN 算法,即使您在索引上有数十亿项,find_neighbors() 方法也只需几毫秒即可提取类似项。Vector Search 还支持自动扩缩,该功能可以根据工作负载的需求自动调整节点数量。

清理

如果您使用的是自己的 Cloud 项目,而不是 Qwiklabs 上的临时项目,请务必在完成本教程后删除所有索引、索引端点和 Cloud Storage 存储桶。否则,剩余的资源可能会产生意外费用。

如果您使用了 Workbench,则可能还需要从控制台中删除笔记本。


# wait for a confirmation
input("Press Enter to delete Index Endpoint, Index and Cloud Storage bucket:")

# delete Index Endpoint
my_index_endpoint.undeploy_all()
my_index_endpoint.delete(force = True)

# delete Index
my_index.delete()

# delete Cloud Storage bucket
! gcloud storage rm {BUCKET_URI} --recursive

实用程序

创建或部署索引可能需要一些时间,届时您可能会中断与 Colab 运行时的连接。如果连接中断,您可以检查 Vector SearchGoogle Cloud 控制台,并使用现有索引继续操作,而不是再次创建或部署新索引。

获取现有索引

如需获取已存在的索引对象,请将以下 your-index-id 替换为索引 ID 并运行该单元。您可以通过检查 Vector Search Google Cloud 控制台来获取索引 ID。在 Google Cloud 控制台的 Agent Platform 部分中,前往部署和使用部分。选择索引

请参阅“索引”

my_index_id = "[your-index-id]"
my_index = aiplatform.MatchingEngineIndex(my_index_id)

获取现有索引端点

如需获取已存在的索引端点对象,请将以下 your-index-endpoint-id 替换为索引端点 ID 并运行该单元。您可以通过检查 Vector Search Google Cloud 控制台来获取索引端点。在 Google Cloud 控制台的 Agent Platform 部分中,前往部署和使用部分。选择将端点编入索引

请参阅“索引端点”

my_index_endpoint_id = "[your-index-endpoint-id]"
my_index_endpoint = aiplatform.MatchingEngineIndexEndpoint(my_index_endpoint_id)