跳到主要内容
Open In ColabOpen on GitHub

谷歌 Vertex AI 特征存储

Google Cloud Vertex 特征库简化了您的机器学习特征管理和在线服务流程,通过让您以低延迟服务您在 Google Cloud BigQuery 中的数据,包括为嵌入执行近似邻居检索的能力

本教程将向您展示如何轻松地直接从您的 BigQuery 数据执行低延迟向量搜索和近似最近邻检索,从而以最少的设置启用强大的机器学习应用。我们将使用 VertexFSVectorStore 类来实现这一点。

此类是 Google Cloud 中能够提供统一数据存储和灵活向量搜索的两类中的一部分

  • BigQuery 向量搜索:使用 BigQueryVectorStore 类,它非常适合无需基础设施设置的快速原型开发和批量检索。
  • 特征库在线存储:使用 VertexFSVectorStore 类,可实现低延迟检索,支持手动或定期数据同步。非常适合面向用户的生产就绪型生成式 AI 应用。

Diagram BQ-VertexFS

开始使用

安装库

%pip install --upgrade --quiet  langchain langchain-google-vertexai "langchain-google-community[featurestore]"

要在当前 Jupyter 运行时中使用新安装的包,您必须重启运行时。您可以通过运行下面的单元格来实现,这会重启当前内核。

import IPython

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

开始之前

设置您的项目 ID

如果您不知道您的项目 ID,请尝试以下操作:

  • 运行 gcloud config list
  • 运行 gcloud projects list
  • 请参阅支持页面:查找项目 ID
PROJECT_ID = ""  # @param {type:"string"}

# Set the project id
! gcloud config set project {PROJECT_ID}

设置区域

您还可以更改 BigQuery 使用的 REGION 变量。了解更多关于 BigQuery 区域的信息。

REGION = "us-central1"  # @param {type: "string"}

设置数据集和表名

它们将是您的 BigQuery 向量存储。

DATASET = "my_langchain_dataset"  # @param {type: "string"}
TABLE = "doc_and_vectors" # @param {type: "string"}

认证您的 Notebook 环境

  • 如果您使用 Colab 运行此 Notebook,请取消注释下面的单元格并继续。
  • 如果您使用 Vertex AI Workbench,请查看此处提供的设置说明。
# from google.colab import auth as google_auth

# google_auth.authenticate_user()

演示:VertexFSVectorStore

创建嵌入类实例

您可能需要通过运行 gcloud services enable aiplatform.googleapis.com --project {PROJECT_ID}(将 {PROJECT_ID} 替换为您的项目名称)来启用项目中的 Vertex AI API。

您可以使用任何 LangChain 嵌入模型

from langchain_google_vertexai import VertexAIEmbeddings

embedding = VertexAIEmbeddings(
model_name="textembedding-gecko@latest", project=PROJECT_ID
)

初始化 VertexFSVectorStore

如果 BigQuery 数据集和表不存在,它们将自动创建。请参阅此处提供的类定义,了解所有可选参数。

from langchain_google_community import VertexFSVectorStore

store = VertexFSVectorStore(
project_id=PROJECT_ID,
dataset_name=DATASET,
table_name=TABLE,
location=REGION,
embedding=embedding,
)

添加文本

注意:由于特征在线存储的创建,首次同步过程大约需要 20 分钟。

all_texts = ["Apples and oranges", "Cars and airplanes", "Pineapple", "Train", "Banana"]
metadatas = [{"len": len(t)} for t in all_texts]

store.add_texts(all_texts, metadatas=metadatas)

您也可以通过执行 sync_data 方法按需启动同步。

store.sync_data()

在生产环境中,您还可以使用 cron_schedule 类参数来设置自动定时同步。例如

store = VertexFSVectorStore(cron_schedule="TZ=America/Los_Angeles 00 13 11 8 *", ...)

搜索文档

query = "I'd like a fruit."
docs = store.similarity_search(query)
print(docs)

按向量搜索文档

query_vector = embedding.embed_query(query)
docs = store.similarity_search_by_vector(query_vector, k=2)
print(docs)

使用元数据过滤器搜索文档

# This should only return "Banana" document.
docs = store.similarity_search_by_vector(query_vector, filter={"len": 6})
print(docs)

添加带有嵌入的文本

您还可以使用 add_texts_with_embeddings 方法引入您自己的嵌入。这对于在生成嵌入之前可能需要自定义预处理的多模态数据特别有用。

items = ["some text"]
embs = embedding.embed(items)

ids = store.add_texts_with_embeddings(
texts=["some text"], embs=embs, metadatas=[{"len": 1}]
)

使用 BigQuery 进行批量服务

您只需使用 .to_bq_vector_store() 方法即可获取 BigQueryVectorStore 对象,该对象为批量用例提供优化的性能。所有强制参数将自动从现有类中转移。请参阅类定义,了解您可以使用的所有参数。

使用 .to_vertex_fs_vector_store() 方法返回 BigQueryVectorStore 同样简单。

store.to_bq_vector_store()  # pass optional VertexFSVectorStore parameters as arguments