从微调到生产：基于 Dataflow 的可扩展嵌入流水线

2025年9月4日

Danny McCormick Staff Software Engineer Google Cloud Dataflow

Ian Ballantyne Senior Developer Relations Engineer Google DeepMind

AI 发展日新月异，而嵌入正成为语义搜索和检索增强生成 (RAG) 等众多现代应用的核心技术。今天，我们很高兴与大家探讨如何利用 Google 全新推出的高效开源嵌入模型 EmbeddingGemma（参数量仅 3.08 亿）。由于其体积小巧，EmbeddingGemma 非常适合用于设备端应用；同时，其高效的性能也为云端赋予了更多可能性，在通过微调进行定制化方面表现尤为突出。我们将展示如何结合 Google Cloud 的 Dataflow 和像 AlloyDB 这样的向量数据库来构建可扩展的实时知识提取流水线。

嵌入和 Dataflow 的强大功能

嵌入是数据的数值向量表征，能够捕捉词语与概念之间的潜在关系，是那些需要从更深层次、概念层面理解信息的应用程序的基础，例如根据查询结果查找语义相似的文档，或者为 RAG 系统中的大型语言模型 (LLM) 提供相关上下文。

为支持这些应用，您需要一个稳健的知识提取流水线，该流水线能够处理非结构化数据、将其转换为向量，并将其加载到专门的向量数据库中。这就是 Dataflow 能够发挥作用的地方，它能够将这些步骤封装到一个代管式流水线中。

在您的流水线中使用像 EmbeddingGemma 这样小巧且高效的开放式模型作为核心，能够使整个过程实现自给自足。这样一来，就可以简化管理，因为无需通过外部网络调用其他服务来完成嵌入步骤。由于这是一个开放模型，所以它可以完全在 Dataflow 中运行。这为安全处理大规模的私有数据集提供了保障。

除了这些操作上的优势之外，EmbeddingGemma 还支持微调，这使您能够根据自身特定的数据嵌入需求对其进行个性化设置；您可以点击此处，了解微调示例。质量与可扩展性同样重要，而 EmbeddingGemma 在这方面也表现出色。在大规模文本嵌入基准 (MTEB) 多语言排行榜上，EmbeddingGemma 是参数量低于 5 亿的排名最高的纯文本多语言嵌入模型。

Dataflow 是全代管式可自动扩展平台，适用于统一的批处理和流式数据处理。通过将像 EmbeddingGemma 这样的模型直接嵌入到 Dataflow 流程中，您将获得多项优势：

数据局部性带来的效率提升：处理工作在 Dataflow 工作线程上进行，无需通过远程过程调用 (RPC) 与独立的推理服务进行交互，也避免了因配额和对多个系统进行自动扩展而产生的问题。您的整个工作流可以整合到一组单一的工作线程中，从而减少资源占用。

统一的系统：一个单一系统即可实现自动扩展、监控和监测功能，从而简化您的运营开销。

可扩展性和简洁性：Dataflow 能够根据需求自动调整您的流水线的规模，并且 Apache Beam 的转换操作能够减少冗余代码。

使用 Dataflow ML 构建提取流水线

典型的知识提取流水线包含四个阶段：从数据源读取数据、对数据进行预处理、生成嵌入以及将数据写入向量数据库。

借助 Dataflow 的 MLTransform（一种功能强大的用于数据准备的 PTransform），整个工作流只需用几行代码就能实现。

使用 MLTransform 生成 Gemma 嵌入

让我们来详细了解一下如何使用新的 Gemma 模型生成文本嵌入。本示例改编自 EmbeddingGemma 笔记本，展示了如何配置 MLTransform 以使用 Hugging Face 模型，然后将结果写入 AlloyDB 中。在 AlloyDB 中，这些嵌入可用于语义搜索。AlloyDB 这样的数据库允许我们将语义搜索与结构化搜索相结合，从而提供高质量且相关的搜索结果。

首先，我们需要定义用于嵌入的模型的名称，并设置一个转换规则，指定要进行嵌入的列以及要使用的模型类型。

import tempfile
import apache_beam as beam
from apache_beam.ml.transforms.base import MLTransform
from apache_beam.ml.transforms.embeddings.huggingface import SentenceTransformerEmbeddings
 
# The new Gemma model for generating embeddings. You can replace this with your fine tuned model just by changing this path.
text_embedding_model_name = 'google/embeddinggemma-300m'
 
# Define the embedding transform with our Gemma model
embedding_transform = SentenceTransformerEmbeddings(
    model_name=text_embedding_model_name, columns=['x']
)

Python

生成嵌入后，我们会直接将输出传输到接收器（通常是一个向量数据库）。为了写入这些嵌入，我们将定义一个基于配置的 VectorDatabaseWriteTransform。

在这个示例中，我们将通过传入 AlloyDBVectorWriterConfig 对象，把 AlloyDB 用作接收器。仅借助这些配置对象，Dataflow 即可支持将嵌入写入多种向量数据库，包括 AlloyDB、CloudSQL 和 BigQuery。

# Define the config used to write to AlloyDB
alloydb_writer_config = AlloyDBVectorWriterConfig(
    connection_config=connection_config,
    table_name=table_name
)
 
# Build and run the pipeline
with beam.Pipeline() as pipeline:
  _ = (
      pipeline
      | "CreateData" >> beam.Create(content) # In production could be replaced by a transform to read from any source
      # MLTransform generates the embeddings
      | "Generate Embeddings" >> MLTransform(
          write_artifact_location=tempfile.mkdtemp()
      ).with_transform(embedding_transform)
      # The output is written to our vector database
      | 'Write to AlloyDB' >> VectorDatabaseWriteTransform(alloydb_writer_config)
  )

Python

通过使用这种简单而强大的方法，您可以并行处理海量数据，使用拥有 3.08 亿参数的 EmbeddingGemma 模型生成嵌入并将结果写入向量数据库，所有这些操作都可以在单个可扩展且经济高效的代管式流水线中完成。

立即开始使用

通过结合使用最新的 Gemma 模型与 Dataflow 的可扩展性和 AlloyDB 等向量数据库的强大向量搜索功能，您可以轻松构建功能强大的新一代 AI 应用。

如要了解详情，请查阅 Dataflow ML 文档，重点关注数据准备和生成嵌入的相关文档。您还可以按照此笔记本中的步骤尝试使用 EmbeddingGemma 构建一个简单的流水线。

对于大型服务器端应用：通过 Gemini API 探索我们最先进的 Gemini 嵌入模型，体验卓越性能和强大的处理能力。

如需详细了解 EmbeddingGemma，请前往 Google 开发者博客阅读我们的发布公告。

发布于：

下一步