Skip to main content
Open In ColabOpen on GitHub

Google Vertex AI 特征存储

Google Cloud Vertex Feature Store 通过让您以低延迟提供数据,简化了您的机器学习特征管理和在线服务流程,数据来源包括 Google Cloud BigQuery,并支持对嵌入向量执行近似邻居检索。

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

这个类是一组两个类的一部分,能够提供统一的数据存储和在 Google Cloud 中的灵活向量搜索:

  • BigQuery向量搜索: 通过BigQueryVectorStore类实现,这使得快速原型设计和无需基础设施设置的批量检索变得理想。
  • Feature Store 在线商店: 通过 VertexFSVectorStore 类别,实现低延迟检索,并支持手动或计划的数据同步。适用于生产级别的面向用户的生成式AI应用程序。

Diagram BQ-VertexFS

Getting Started

安装库

%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"}

认证您的笔记本环境

  • 如果您正在使用Colab运行此笔记本,请取消注释下方的单元格并继续。
  • 如果您正在使用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
)
API 参考:VertexAI 嵌入

初始化 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