Google 发布 EmbeddingGemma 高效多语言嵌入模型
Welcome EmbeddingGemma, Google's new efficient embedding model
Google 发布 EmbeddingGemma,一个面向端侧的多语言嵌入模型,参数量 308M,支持 2K 上下文窗口和 100 多种语言,量化后 RAM 占用低于 200 MB。
原文给出模型架构、MRL 截断维度和多框架用法,读者可以据此评估它在移动端 RAG 和检索场景的可行性。
TL;DR
今天,Google 发布 EmbeddingGemma,这是一个最先进的多语言嵌入模型,非常适合端侧使用场景。该模型专为速度和效率而设计,具有紧凑的 308M 参数规模和 2K 上下文窗口,为移动端 RAG 流水线、智能体等解锁了新的可能性。EmbeddingGemma 经过训练可支持超过 100 种语言,并且在撰写本文时,是 Massive Text Embedding Benchmark(MTEB)上排名最高的 5 亿参数以下的纯文本多语言嵌入模型。
目录
简介
文本嵌入已成为现代自然语言应用的支柱,将单词、句子和文档转化为能够捕捉含义、情感和意图的稠密向量。这些向量支持在庞大的语料库中进行快速相似度搜索、聚类、分类和检索,为从推荐引擎和语义搜索到检索增强生成和代码搜索工具的一切提供支持。计算这些嵌入的嵌入模型被广泛使用,在 Hugging Face 上每月下载量远超 2 亿次。
在此基础上,Google DeepMind 的 EmbeddingGemma 作为迄今最新、能力最强的小型多语言嵌入模型问世。EmbeddingGemma 仅具有 308M 参数、2k token 上下文窗口,并支持超过 100 种语言,在 Massive Multilingual Text Embedding Benchmark(MMTEB)上实现了最先进的性能,同时在量化后内存占用保持在 200 MB 以下。
这些设计选择造就了一个非常实用的开源工具,可在日常设备上计算高质量的多语言嵌入。
在这篇博客文章中,我们介绍 EmbeddingGemma 的架构和训练,并向你展示如何将该模型与 Sentence Transformers、LangChain、LlamaIndex、Haystack、txtai、Transformers.js、Text Embedding Inference 和 ONNX 等各种框架一起使用。
随后,我们演示如何针对你的领域微调 EmbeddingGemma,以获得更强的性能。在我们的示例中,我们在 Medical Instruction and Retrieval Dataset(MIRIAD)上微调 EmbeddingGemma。得到的模型 sentence-transformers/embeddinggemma-300m-medical 在我们的任务上实现了最先进的性能:根据详细的医学问题检索科学医学论文段落。它甚至在该任务上超越了规模两倍的模型。
架构
EmbeddingGemma 基于 Gemma3 transformers 主干构建,但经过修改,使用双向注意力而非因果(单向)注意力。这意味着序列中较早的 token 可以关注较晚的 token,从而有效地将架构从解码器转变为编码器。在检索等嵌入任务上,编码器模型可以超越作为解码器的 LLM(Weller et al., 2025)。借助这一主干,模型可以一次性处理多达 2048 个 token,足以应对典型的检索输入,尤其是考虑到更大的输入往往会导致文本嵌入中的信息丢失。
除了基于 Gemma3 的新编码器主干(用于生成 token 嵌入)之外,一个平均池化层会将这些 token 嵌入转换为文本嵌入。最后,两个全连接层将文本嵌入转换为最终形式,即一个 768 维向量。
EmbeddingGemma 模型使用 Matryoshka Representation Learning (MRL) 进行训练,允许你按需将 768 维输出截断为 512、256 或 128 维。这会加快下游处理速度,并降低内存和磁盘空间占用。有关如何执行此截断的代码片段,请参阅 Sentence Transformers usage。
该模型使用精心策划的多语言语料库进行训练,总计约 3200 亿个 token。该专有数据集融合了公开可用的网络文本、代码和技术文档,以及合成的任务特定示例。它已经过过滤,以避免儿童性虐待材料(CSAM)、敏感数据以及低质量或不安全的内容。
评估
EmbeddingGemma 在 MMTEB(多语言,v2)和 MTEB(英语,v2)套件上进行了基准测试,这些套件涵盖广泛的任务、领域和语言。尽管其参数规模仅为 308M,该模型仍持续优于可比的基线模型,同时保持极小的内存占用。
| MTEB(多语言,v2)性能 | MTEB(英语,v2)性能 |
|---|---|
|
|
结果将列在官方 MTEB Leaderboard 上。我们排除任何在超过 20% 的 MTEB 数据上训练过的模型,以减轻潜在的过拟合。
演示
该 demo 也可以 全屏体验。
在桌面设备上亲自体验 demo。
用法
EmbeddingGemma 已与许多流行工具集成,使其易于融入你现有的工作流和应用程序。该模型已集成到 Sentence Transformers 中,因此也集成到在幕后使用 Sentence Transformers 的项目中,例如 LangChain、LlamaIndex、Haystack 和 txtai。请参阅下面的示例,以开始使用你偏好的框架。
对于生产部署,你可以使用 Text Embeddings Inference(TEI)在各种硬件配置上高效地提供模型服务,并且你可以使用 Transformers.js 用于 Web 应用程序。
无论你选择哪种框架,你都应注意 prompts。对于嵌入模型,prompts 会被前置到输入文本中,以便模型区分不同的任务。EmbeddingGemma 是使用这些 prompt 名称和 prompts 训练的,因此在使用该模型时也应包含它们:
query:"task: search result | query: ",document:"title: none | text: ",BitextMining:"task: search result | query: ",Clustering:"task: clustering | query: ",Classification:"task: classification | query: ",InstructionRetrieval:"task: code retrieval | query: ",MultilabelClassification:"task: classification | query: ",PairClassification:"task: sentence similarity | query: ",Reranking:"task: search result | query: ",Retrieval-query:"task: search result | query: ",Retrieval-document:"title: none | text: ",STS:"task: sentence similarity | query: ",Summarization:"task: summarization | query: "
在 Sentence Transformers 中,调用 model.encode_query 和 model.encode_document 时会自动使用 query 和 document 提示词,但对于其他框架,你可能需要:$
- 指定提示词名称(例如 "Reranking"),
- 指定提示词字符串(例如 "task: search result | query: "),或
- 手动将提示词添加到输入文本前面。
以下示例脚本将使用各种框架演示这一点。
Sentence Transformers
你需要安装以下软件包:
pip install git+https://github.com/huggingface/transformers@v4.56.0-Embedding-Gemma-preview
pip install sentence-transformers>=5.0.0
检索
使用 Sentence Transformers 进行推理相当简单,请参阅这个语义搜索示例:
from sentence_transformers import SentenceTransformer
# Download from the 🤗 Hub
model = SentenceTransformer("google/embeddinggemma-300m")
# Run inference with queries and documents
query = "Which planet is known as the Red Planet?"
documents = [
"Venus is often called Earth's twin because of its similar size and proximity.",
"Mars, known for its reddish appearance, is often referred to as the Red Planet.",
"Jupiter, the largest planet in our solar system, has a prominent red spot.",
"Saturn, famous for its rings, is sometimes mistaken for the Red Planet."
]
query_embeddings = model.encode_query(query)
document_embeddings = model.encode_document(documents)
print(query_embeddings.shape, document_embeddings.shape)
# (768,) (4, 768)
# Compute similarities to determine a ranking
similarities = model.similarity(query_embeddings, document_embeddings)
print(similarities)
# tensor([[0.3011, 0.6359, 0.4930, 0.4889]])
# Convert similarities to a ranking
ranking = similarities.argsort(descending=True)[0]
print(ranking)
# tensor([1, 2, 3, 0])
- Sentence Transformers
encode_query方法文档 - Sentence Transformers
encode_document方法文档 - Sentence Transformers
similarity方法文档
点击查看非检索代码
如果你不打算将此模型用于信息检索,那么你最好使用最通用的 encode 方法,并搭配以下选项中能最好地描述你下游任务的模型提示词:
BitextMining:查找两种语言中的翻译句子对。Clustering:查找相似文本以将它们分组。Classification:为文本分配预定义标签。InstructionRetrieval:根据自然语言指令检索相关代码片段。MultilabelClassification:为文本分配多个标签。PairClassification:为文本分配预定义标签。Reranking:根据相关性重新排列搜索结果。Retrieval-query:根据查询检索文档。Retrieval-document:根据其内容检索文档。STS:计算文本之间的语义文本相似度。Summarization:生成文本的简洁摘要。
from sentence_transformers import SentenceTransformer
# Download from the 🤗 Hub
model = SentenceTransformer("google/embeddinggemma-300m")
# Let's inspect the configured prompts
print(model.prompts)
# {
# "query": "task: search result | query: ",
# "document": "title: none | text: ",
# "BitextMining": "task: search result | query: ",
# "Clustering": "task: clustering | query: ",
# "Classification": "task: classification | query: ",
# "InstructionRetrieval": "task: code retrieval | query: ",
# "MultilabelClassification": "task: classification | query: ",
# "PairClassification": "task: sentence similarity | query: ",
# "Reranking": "task: search result | query: ",
# "Retrieval-query": "task: search result | query: ",
# "Retrieval-document": "title: none | text: ",
# "STS": "task: sentence similarity | query: ",
# "Summarization": "task: summarization | query: ",
# }
# Compute semantic textual similarity using texts, so let's use the STS prompt
texts = [
"The weather is beautiful today.",
"It's a lovely day outside.",
"The stock market crashed yesterday.",
"I enjoy programming with Python."
]
embeddings = model.encode(texts, prompt_name="STS")
print(embeddings.shape)
# (4, 768)
# Compute similarities
similarities = model.similarity(embeddings, embeddings)
print(similarities)
"""
tensor([[1.0000, 0.9305, 0.4660, 0.4326],
[0.9305, 1.0000, 0.4227, 0.4434],
[0.4660, 0.4227, 1.0000, 0.2638],
[0.4326, 0.4434, 0.2638, 1.0000]])
"""
点击查看如何截断嵌入维度以实现更快、更便宜的搜索
由于 google/embeddinggemma-300m 是使用 MRL 训练的,该模型生成的嵌入可以截断为较低的维度,而不会显著损害评估性能。较低维度的嵌入在磁盘和内存中的存储成本都更低,并且在检索、聚类或分类等下游任务中速度更快。
在 Sentence Transformers 中,你可以在 SentenceTransformer 初始化时或调用 model.encode/model.encode_query/model.encode_document 时使用 truncate_dim 参数设置较低的维度:
from sentence_transformers import SentenceTransformer
# Download from the 🤗 Hub
model = SentenceTransformer("google/embeddinggemma-300m", truncate_dim=256)
# Run inference with queries and documents
query = "Which planet is known as the Red Planet?"
documents = [
"Venus is often called Earth's twin because of its similar size and proximity.",
"Mars, known for its reddish appearance, is often referred to as the Red Planet.",
"Jupiter, the largest planet in our solar system, has a prominent red spot.",
"Saturn, famous for its rings, is sometimes mistaken for the Red Planet."
]
query_embeddings = model.encode_query(query)
document_embeddings = model.encode_document(documents)
print(query_embeddings.shape, document_embeddings.shape)
# (256,) (4, 256)
# Compute similarities to determine a ranking
similarities = model.similarity(query_embeddings, document_embeddings)
print(similarities)
# tensor([[0.4016, 0.6715, 0.5283, 0.5261]])
# Convert similarities to a ranking
ranking = similarities.argsort(descending=True)[0]
print(ranking)
# tensor([1, 2, 3, 0])
请注意,尽管使用的嵌入比全尺寸嵌入小 3 倍,但排名仍然保持不变。
LangChain
如果你愿意,也可以使用 LangChain HuggingFaceEmbeddings,它在底层使用 Sentence Transformers。请注意,你需要告诉 LangChain 分别对查询和文档使用名为 "query" 和 "document" 的提示词。此示例涉及一个简单的信息检索设置,但相同的嵌入模型也可用于更复杂的场景。
你需要安装以下软件包:
pip install git+https://github.com/huggingface/transformers@v4.56.0-Embedding-Gemma-preview
pip install sentence-transformers
pip install langchain
pip install langchain-community
pip install langchain-huggingface
pip install faiss-cpu
from langchain.docstore.document import Document
from langchain_community.vectorstores import FAISS
from langchain_huggingface.embeddings import HuggingFaceEmbeddings
# Download the model from the 🤗 Hub. Also specify to use the "query" and "document" prompts
# as defined in the model configuration, as LangChain doesn't automatically use them.
# See https://huggingface.co/google/embeddinggemma-300m/blob/main/config_sentence_transformers.json
embedder = HuggingFaceEmbeddings(
model_name="google/embeddinggemma-300m",
query_encode_kwargs={"prompt_name": "query"},
encode_kwargs={"prompt_name": "document"}
)
data = [
"Venus is often called Earth's twin because of its similar size and proximity.",
"Mars, known for its reddish appearance, is often referred to as the Red Planet.",
"Jupiter, the largest planet in our solar system, has a prominent red spot.",
"Saturn, famous for its rings, is sometimes mistaken for the Red Planet."
]
# Create documents for the vector store
documents = [Document(page_content=text, metadata={"id": i}) for i, text in enumerate(data)]
# Create vector store using FAISS. Setting distance_strategy to "MAX_INNER_PRODUCT" uses
# FAISS' FlatIndexIP behind the scenes, which is optimized for inner product search. This
# is what the model was trained for
vector_store = FAISS.from_documents(documents, embedder, distance_strategy="MAX_INNER_PRODUCT")
# Search for top 3 similar documents
query = "Which planet is known as the Red Planet?"
results = vector_store.similarity_search_with_score(query, k=3)
# Print results
for doc, score in results:
print(f"Text: {doc.page_content} (score: {score:.4f})")
"""
Text: Mars, known for its reddish appearance, is often referred to as the Red Planet. (score: 0.6359)
Text: Jupiter, the largest planet in our solar system, has a prominent red spot. (score: 0.4930)
Text: Saturn, famous for its rings, is sometimes mistaken for the Red Planet. (score: 0.4889)
"""
LlamaIndex
EmbeddingGemma 在 LlamaIndex 中也受支持,因为它在底层使用 Sentence Transformers。为了获得正确的行为,你需要指定模型配置中定义的查询和文档提示词。否则,你的性能将无法达到最优。此脚本展示了一个使用 EmbeddingGemma 与 LlamaIndex 的基础示例,但你也可以在更复杂的场景中使用 HuggingFaceEmbedding 类。
你需要安装以下软件包:
pip install git+https://github.com/huggingface/transformers@v4.56.0-Embedding-Gemma-preview
pip install sentence-transformers
pip install llama-index
pip install llama-index-embeddings-huggingface
pip install llama-index-vector-stores-faiss
import faiss
from llama_index.core.schema import TextNode
from llama_index.core.vector_stores import VectorStoreQuery
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from llama_index.vector_stores.faiss import FaissVectorStore
# Download from the 🤗 Hub. Also specify the query and document prompts as
# defined in the model configuration, as LlamaIndex doesn't automatically load them.
# See https://huggingface.co/google/embeddinggemma-300m/blob/main/config_sentence_transformers.json
embeddings = HuggingFaceEmbedding(
model_name="google/embeddinggemma-300m",
query_instruction="task: search result | query: ",
text_instruction="title: none | text: ",
)
data = [
"Venus is often called Earth's twin because of its similar size and proximity.",
"Mars, known for its reddish appearance, is often referred to as the Red Planet.",
"Jupiter, the largest planet in our solar system, has a prominent red spot.",
"Saturn, famous for its rings, is sometimes mistaken for the Red Planet."
]
# Create a sample vector store
store = FaissVectorStore(faiss_index=faiss.IndexFlatIP(768))
store.add([TextNode(id=i, text=text, embedding=embeddings.get_text_embedding(text)) for i, text in enumerate(data)])
# Search for top k similar documents
query = "Which planet is known as the Red Planet?"
query_embedding = embeddings.get_query_embedding(query)
results = store.query(VectorStoreQuery(query_embedding=query_embedding, similarity_top_k=3))
# Print results
for idx, score in zip(results.ids, results.similarities):
print(f"Text: {data[int(idx)]} (score: {score:.4f})")
"""
Text: Mars, known for its reddish appearance, is often referred to as the Red Planet. (score: 0.6359)
Text: Jupiter, the largest planet in our solar system, has a prominent red spot. (score: 0.4930)
Text: Saturn, famous for its rings, is sometimes mistaken for the Red Planet. (score: 0.4889)
"""
Haystack
EmbeddingGemma 也可以与 Haystack 一起使用,Haystack 是一个用于构建生产级搜索和语言应用的框架。与 LangChain 和 LlamaIndex 一样,Haystack 在底层使用 Sentence Transformers,并且需要你指定适当的提示词。以下示例展示了如何使用 EmbeddingGemma 与 Haystack 搭建一个基本的检索流水线。
你需要安装以下软件包:
pip install git+https://github.com/huggingface/transformers@v4.56.0-Embedding-Gemma-preview
pip install sentence-transformers
pip install haystack-ai
from haystack import Document, Pipeline
from haystack.components.embedders import SentenceTransformersDocumentEmbedder, SentenceTransformersTextEmbedder
from haystack.components.retrievers import InMemoryEmbeddingRetriever
from haystack.document_stores.in_memory import InMemoryDocumentStore
# Initialize the document store
document_store = InMemoryDocumentStore()
# Initialize the document and query embedders
document_embedder = SentenceTransformersDocumentEmbedder(
model="google/embeddinggemma-300m", encode_kwargs={"prompt_name": "document"}
)
query_embedder = SentenceTransformersTextEmbedder(
model="google/embeddinggemma-300m", encode_kwargs={"prompt_name": "query"}
)
document_embedder.warm_up()
query_embedder.warm_up()
data = [
"Venus is often called Earth's twin because of its similar size and proximity.",
"Mars, known for its reddish appearance, is often referred to as the Red Planet.",
"Jupiter, the largest planet in our solar system, has a prominent red spot.",
"Saturn, famous for its rings, is sometimes mistaken for the Red Planet.",
]
# Convert to Haystack documents and write to document store
documents = [Document(content=text, id=str(i)) for i, text in enumerate(data)]
documents_with_embeddings = document_embedder.run(documents=documents)["documents"]
document_store.write_documents(documents_with_embeddings)
# Create a query pipeline using a query embedder and compatible retriever
query_pipeline = Pipeline()
query_pipeline.add_component("text_embedder", query_embedder)
query_pipeline.add_component("retriever", InMemoryEmbeddingRetriever(document_store=document_store, top_k=3))
query_pipeline.connect("text_embedder.embedding", "retriever.query_embedding")
# Search for top 3 similar documents
query = "Which planet is known as the Red Planet?"
results = query_pipeline.run({"text_embedder": {"text": query}})
# Print results
for document in results["retriever"]["documents"]:
print(f"Text: {document.content} (score: {document.score:.4f})")
"""
Text: Mars, known for its reddish appearance, is often referred to as the Red Planet. (score: 0.6359)
Text: Jupiter, the largest planet in our solar system, has a prominent red spot. (score: 0.4930)
Text: Saturn, famous for its rings, is sometimes mistaken for the Red Planet. (score: 0.4889)
"""
txtai
txtai 也与 EmbeddingGemma 兼容。与其他框架一样,txtai 在底层使用 Sentence Transformers,并且需要适当的提示词才能让 EmbeddingGemma 发挥最佳性能。以下示例演示了如何使用 txtai 搭建一个基本的检索系统。
你需要安装以下软件包:
pip install git+https://github.com/huggingface/transformers@v4.56.0-Embedding-Gemma-preview
pip install sentence-transformers
pip install txtai
from txtai import Embeddings
# Download from the 🤗 Hub. Also specify the query and document prompts as
# defined in the model configuration, as txtai doesn't automatically load them.
# See https://huggingface.co/google/embeddinggemma-300m/blob/main/config_sentence_transformers.json
embeddings = Embeddings(
path="google/embeddinggemma-300m",
method="sentence-transformers",
instructions={
"query": "task: search result | query: ",
"data": "title: none | text: ",
}
)
data = [
"Venus is often called Earth's twin because of its similar size and proximity.",
"Mars, known for its reddish appearance, is often referred to as the Red Planet.",
"Jupiter, the largest planet in our solar system, has a prominent red spot.",
"Saturn, famous for its rings, is sometimes mistaken for the Red Planet."
]
# Create a sample vector store
embeddings.index(data)
# Search for top k similar documents
query = "Which planet is known as the Red Planet?"
results = embeddings.search(query, 3)
# Print results
for idx, score in results:
print(f"Text: {data[int(idx)]} (score: {score:.4f})")
"""
Text: Mars, known for its reddish appearance, is often referred to as the Red Planet. (score: 0.6359)
Text: Jupiter, the largest planet in our solar system, has a prominent red spot. (score: 0.4930)
Text: Saturn, famous for its rings, is sometimes mistaken for the Red Planet. (score: 0.4889)
"""
Transformers.js
你甚至可以使用 Transformers.js 在浏览器中 100% 本地运行 EmbeddingGemma!如果你还没有安装,可以通过以下命令从 NPM 安装该库:
npm i @huggingface/transformers
然后你可以按如下方式计算嵌入:
import { AutoModel, AutoTokenizer, matmul } from "@huggingface/transformers";
// Download from the 🤗 Hub
const model_id = "onnx-community/embeddinggemma-300m-ONNX";
const tokenizer = await AutoTokenizer.from_pretrained(model_id);
const model = await AutoModel.from_pretrained(model_id, {
dtype: "fp32", // Options: "fp32" | "q8" | "q4"
});
// Run inference with queries and documents
const prefixes = {
query: "task: search result | query: ",
document: "title: none | text: ",
};
const query = prefixes.query + "Which planet is known as the Red Planet?";
const documents = [
"Venus is often called Earth's twin because of its similar size and proximity.",
"Mars, known for its reddish appearance, is often referred to as the Red Planet.",
"Jupiter, the largest planet in our solar system, has a prominent red spot.",
"Saturn, famous for its rings, is sometimes mistaken for the Red Planet.",
].map((x) => prefixes.document + x);
const inputs = await tokenizer([query, ...documents], { padding: true });
const { sentence_embedding } = await model(inputs);
// Compute similarities to determine a ranking
const scores = await matmul(sentence_embedding, sentence_embedding.transpose(1, 0));
const similarities = scores.tolist()[0].slice(1);
console.log(similarities);
// [ 0.30109718441963196, 0.6358831524848938, 0.4930494725704193, 0.48887503147125244 ]
// Convert similarities to a ranking
const ranking = similarities.map((score, index) => ({ index, score })).sort((a, b) => b.score - a.score);
console.log(ranking);
// [
// { index: 1, score: 0.6358831524848938 },
// { index: 2, score: 0.4930494725704193 },
// { index: 3, score: 0.48887503147125244 },
// { index: 0, score: 0.30109718441963196 }
// ]
Text Embeddings Inference
你可以使用 Text Embeddings Inference (TEI) 1.8.1 或更高版本,轻松地将 EmbeddingGemma 部署到开发和生产环境中。
- CPU:
docker run -p 8080:80 ghcr.io/huggingface/text-embeddings-inference:cpu-1.8.1 --model-id google/embeddinggemma-300m --dtype float32
- 使用 ONNX Runtime 的 CPU:
docker run -p 8080:80 ghcr.io/huggingface/text-embeddings-inference:cpu-1.8.1 --model-id onnx-community/embeddinggemma-300m-ONNX --dtype float32 --pooling mean
- NVIDIA CUDA:
docker run --gpus all --shm-size 1g -p 8080:80 ghcr.io/huggingface/text-embeddings-inference:cuda-1.8.1 --model-id google/embeddinggemma-300m --dtype float32
如果你使用
cuda-1.8.1标签运行 Docker 容器,它包含对多种 GPU 架构的支持:Turing、Ampere、Ada Lovelace 和 Hopper。如果你想要一个仅针对你的 GPU 的更轻量镜像,可以改用特定标签,例如turing-1.8.1、1.8.1和86-1.8.1(Ampere)、89-1.8.1(Ada Lovelace)或hopper-1.8.1。
部署完成后,无论使用何种设备或运行时,你都可以利用基于 OpenAI Embeddings API 规范的 /v1/embeddings 端点来生成嵌入。
curl http://0.0.0.0:8080/v1/embeddings -H "Content-Type: application/json" -d '{"model":"google/embeddinggemma-300m","input":["task: search result | query: Which planet is known as the Red Planet?","task: search result | query: Where did Amelia Earhart first fly?"]}'
或者,你也可以利用 Text Embeddings Inference Embeddings API 中的 /embed 端点,它支持 prompt_name 参数,这意味着无需手动将提示词前置到输入中,而是通过 prompt_name 来选择它。
curl http://0.0.0.0:8080/embed -H "Content-Type: application/json" -d '{"inputs":["Which planet is known as the Red Planet?","Where did Amelia Earthart first fly?"],"prompt_name":"query","normalize":true}'
此外,请注意,由于
google/embeddinggemma-300m是使用 Matryoshka Representation Learning (MRL) 训练的,你还可以在/v1/embeddings和/embed上利用dimensions参数,将嵌入截断为更低的维度(512、256 和 128),而不会损害评估性能。
ONNX Runtime
你也可以直接使用 ONNX Runtime 运行该模型,使其具有高度的可移植性和跨平台兼容性。下面的示例展示了在 Python 中的用法,但同样的方法也可以应用于其他语言(Java、C#、C++ 等)。
from huggingface_hub import hf_hub_download
import onnxruntime as ort
from transformers import AutoTokenizer
# Download from the 🤗 Hub
model_id = "onnx-community/embeddinggemma-300m-ONNX"
model_path = hf_hub_download(model_id, subfolder="onnx", filename="model.onnx") # Download graph
hf_hub_download(model_id, subfolder="onnx", filename="model.onnx_data") # Download weights
session = ort.InferenceSession(model_path)
tokenizer = AutoTokenizer.from_pretrained(model_id)
# Run inference with queries and documents
prefixes = {
"query": "task: search result | query: ",
"document": "title: none | text: ",
}
query = prefixes["query"] + "Which planet is known as the Red Planet?"
documents = [
"Venus is often called Earth's twin because of its similar size and proximity.",
"Mars, known for its reddish appearance, is often referred to as the Red Planet.",
"Jupiter, the largest planet in our solar system, has a prominent red spot.",
"Saturn, famous for its rings, is sometimes mistaken for the Red Planet."
]
documents = [prefixes["document"] + x for x in documents]
inputs = tokenizer([query] + documents, padding=True, return_tensors="np")
_, sentence_embedding = session.run(None, inputs.data)
print(sentence_embedding.shape) # (5, 768)
# Compute similarities to determine a ranking
query_embeddings = sentence_embedding[0]
document_embeddings = sentence_embedding[1:]
similarities = query_embeddings @ document_embeddings.T
print(similarities) # [0.30109745 0.635883 0.49304956 0.48887485]
# Convert similarities to a ranking
ranking = similarities.argsort()[::-1]
print(ranking) # [1 2 3 0]
微调
与所有兼容 Sentence Transformers 库的模型一样,EmbeddingGemma 可以轻松地在你特定的数据集上进行微调。为了展示这一点,我们将在 Medical Instruction and RetrIeval Dataset (MIRIAD) 数据集上微调 google/embeddinggemma-300m,这样我们微调后的模型就会特别擅长根据详细的医学问题从科学医学论文中查找最多 1000 个 token 的段落。这些段落可以作为生成模型的关键上下文,从而更有效地回答问题。
下面,你可以使用可展开的标签页探索微调过程的每个关键组件。每个标签页都包含相关代码和详细说明。
模型
from sentence_transformers import SentenceTransformer, SentenceTransformerModelCardData
model = SentenceTransformer(
"google/embeddinggemma-300m",
model_card_data=SentenceTransformerModelCardData(
language="en",
license="apache-2.0",
model_name="EmbeddingGemma-300m trained on the Medical Instruction and RetrIeval Dataset (MIRIAD)",
),
)
# SentenceTransformer(
# (0): Transformer({'max_seq_length': 1024, 'do_lower_case': False, 'architecture': 'Gemma3TextModel'})
# (1): Pooling({'word_embedding_dimension': 768, 'pooling_mode_cls_token': False, 'pooling_mode_mean_tokens': True, 'pooling_mode_max_tokens': False, 'pooling_mode_mean_sqrt_len_tokens': False, 'pooling_mode_weightedmean_tokens': False, 'pooling_mode_lasttoken': False, 'include_prompt': True})
# (2): Dense({'in_features': 768, 'out_features': 3072, 'bias': False, 'activation_function': 'torch.nn.modules.linear.Identity'})
# (3): Dense({'in_features': 3072, 'out_features': 768, 'bias': False, 'activation_function': 'torch.nn.modules.linear.Identity'})
# (4): Normalize()
# )
这段代码从 Hugging Face 加载 EmbeddingGemma 模型,并可选地附带模型卡片元数据,用于文档记录和分享。SentenceTransformer 类加载模型权重和配置,而 model_card_data 参数则附加有助于纳入自动生成的模型卡片的元数据。
数据集
from datasets import load_dataset
train_dataset = load_dataset("tomaarsen/miriad-4.4M-split", split="train").select(range(100_000))
eval_dataset = load_dataset("tomaarsen/miriad-4.4M-split", split="eval").select(range(1_000))
test_dataset = load_dataset("tomaarsen/miriad-4.4M-split", split="test").select(range(1_000))
# Dataset({
# features: ['question', 'passage_text'],
# num_rows: 100000
# })
# Dataset({
# features: ['question', 'passage_text'],
# num_rows: 1000
# })
# Dataset({
# features: ['question', 'passage_text'],
# num_rows: 1000
# })
这段代码加载 MIRIAD 数据集,更准确地说,是加载一个已划分为训练集、评估集和测试集的副本。使用大规模、高质量的数据集可确保模型学习到有意义的表示,而取子集则能加快实验速度。load_dataset 函数从 Hugging Face Datasets 获取数据集,.select() 方法则限制每个划分的样本数量。
损失函数
from sentence_transformers.losses import CachedMultipleNegativesRankingLoss
loss = CachedMultipleNegativesRankingLoss(model, mini_batch_size=8)
这段代码定义了训练用的损失函数,使用 Cached Multiple Negatives Ranking Loss (CMNRL)。CMNRL 对检索任务很有效,因为它利用批内负样本高效地训练模型区分正确与错误的配对。该损失函数接收问答对,并将批次中的其他答案视为负样本,从而最大化嵌入空间中不相关配对之间的距离。mini_batch_size 参数控制内存使用量,但不影响训练动态。
建议将此损失函数与 SentenceTransformerTrainingArguments 中较大的 per_device_train_batch_size 以及 CachedMultipleNegativesRankingLoss 中较低的 mini_batch_size 配合使用,以在低内存占用下获得强训练信号。此外,推荐使用 NO_DUPLICATES 批采样器以避免意外的假负样本。
训练参数
from sentence_transformers.training_args import BatchSamplers
from sentence_transformers import SentenceTransformerTrainingArguments
run_name = "embeddinggemma-300m-medical-100k"
args = SentenceTransformerTrainingArguments(
output_dir=f"models/{run_name}",
num_train_epochs=1,
per_device_train_batch_size=128,
per_device_eval_batch_size=128,
learning_rate=2e-5,
warmup_ratio=0.1,
fp16=True, # Set to False if your GPU can't run FP16
bf16=False, # Set to True if your GPU supports BF16
batch_sampler=BatchSamplers.NO_DUPLICATES,
prompts={
"question": model.prompts["query"],
"passage_text": model.prompts["document"],
},
eval_strategy="steps",
eval_steps=100,
save_strategy="steps",
save_steps=100,
save_total_limit=2,
logging_steps=20,
run_name=run_name,
)
这段代码设置了训练、评估和日志记录的所有超参数和配置。恰当的训练参数对于高效、稳定且可复现的训练至关重要。这些参数控制批次大小、学习率、混合精度、评估和保存频率等。值得注意的是,prompts 字典将数据集列映射到模型用于区分查询与文档的提示。
评估器
from sentence_transformers.evaluation import InformationRetrievalEvaluator
queries = dict(enumerate(eval_dataset["question"]))
corpus = dict(enumerate(eval_dataset["passage_text"] + train_dataset["passage_text"][:30_000]))
relevant_docs = {idx: [idx] for idx in queries}
dev_evaluator = InformationRetrievalEvaluator(
queries=queries,
corpus=corpus,
relevant_docs=relevant_docs,
name="miriad-eval-1kq-31kd",
show_progress_bar=True,
)
dev_evaluator(model)
这段代码设置了一个用于信息检索的评估器,使用查询和语料库来衡量模型性能。训练期间的评估有助于监控进度并避免过拟合。评估器通过检查模型是否为每个查询检索到正确的段落来计算检索指标(NDCG、MRR、Recall、Precision、MAP 等)。它可以在训练前、训练中和训练后运行,结果将被记录并纳入自动生成的模型卡片中。
请注意,此代码片段特别使用了全部(1k)评估问题,针对包含全部(1k)评估段落和 30k 训练段落的语料库,总计 31k 文档进行评估。仅针对评估段落进行评估对模型来说过于简单。
训练器
from sentence_transformers import SentenceTransformerTrainer
trainer = SentenceTransformerTrainer(
model=model,
args=args,
train_dataset=train_dataset,
eval_dataset=eval_dataset,
loss=loss,
evaluator=dev_evaluator,
)
trainer.train()
这段代码初始化并运行训练循环,协调所有组件。
完整微调脚本
以下是完整脚本,整合了上述所有组件:
import logging
import traceback
from datasets import load_dataset
from sentence_transformers import (
SentenceTransformer,
SentenceTransformerModelCardData,
SentenceTransformerTrainer,
SentenceTransformerTrainingArguments,
)
from sentence_transformers.evaluation import InformationRetrievalEvaluator
from sentence_transformers.losses import CachedMultipleNegativesRankingLoss
from sentence_transformers.training_args import BatchSamplers
# Set the log level to INFO to get more information
logging.basicConfig(format="%(asctime)s - %(message)s", datefmt="%Y-%m-%d %H:%M:%S", level=logging.INFO)
# 1. Load a model to finetune with 2. (Optional) model card data
model = SentenceTransformer(
"google/embeddinggemma-300m",
model_card_data=SentenceTransformerModelCardData(
language="en",
license="apache-2.0",
model_name="EmbeddingGemma-300m trained on the Medical Instruction and RetrIeval Dataset (MIRIAD)",
),
)
# 3. Load a dataset to finetune on
train_dataset = load_dataset("tomaarsen/miriad-4.4M-split", split="train").select(range(100_000))
eval_dataset = load_dataset("tomaarsen/miriad-4.4M-split", split="eval").select(range(1_000))
test_dataset = load_dataset("tomaarsen/miriad-4.4M-split", split="test").select(range(1_000))
# 4. Define a loss function. CachedMultipleNegativesRankingLoss (CMNRL) is a special variant of MNRL (a.k.a. InfoNCE),
# which take question-answer pairs (or triplets, etc.) as input. It will take answers from other questions in the batch
# as wrong answers, reducing the distance between the question and the true answer while increasing the distance to the
# wrong answers, in the embedding space.
# The (C)MNRL losses benefit from larger `per_device_train_batch_size` in the Training Arguments, as they can leverage
# more in-batch negative samples. At the same time, the `mini_batch_size` does not affect training performance, but it
# does limit the memory usage. A good trick is setting a high `per_device_train_batch_size` while keeping
# `mini_batch_size` small.
loss = CachedMultipleNegativesRankingLoss(model, mini_batch_size=8)
# 5. (Optional) Specify training arguments
run_name = "embeddinggemma-300m-medical-100k"
args = SentenceTransformerTrainingArguments(
# Required parameter:
output_dir=f"models/{run_name}",
# Optional training parameters:
num_train_epochs=1,
per_device_train_batch_size=128,
per_device_eval_batch_size=128,
learning_rate=2e-5,
warmup_ratio=0.1,
fp16=True, # Set to False if you get an error that your GPU can't run on FP16
bf16=False, # Set to True if you have a GPU that supports BF16
batch_sampler=BatchSamplers.NO_DUPLICATES, # (Cached)MultipleNegativesRankingLoss benefits from no duplicate samples in a batch
prompts={ # Map training column names to model prompts
"question": model.prompts["query"],
"passage_text": model.prompts["document"],
},
# Optional tracking/debugging parameters:
eval_strategy="steps",
eval_steps=100,
save_strategy="steps",
save_steps=100,
save_total_limit=2,
logging_steps=20,
run_name=run_name, # Will be used in W&B if `wandb` is installed
)
# 6. (Optional) Create an evaluator using the evaluation queries and 31k answers & evaluate the base model
queries = dict(enumerate(eval_dataset["question"]))
corpus = dict(enumerate(eval_dataset["passage_text"] + train_dataset["passage_text"][:30_000]))
relevant_docs = {idx: [idx] for idx in queries}
dev_evaluator = InformationRetrievalEvaluator(
queries=queries,
corpus=corpus,
relevant_docs=relevant_docs,
name="miriad-eval-1kq-31kd", # 1k questions, 31k passages
show_progress_bar=True,
)
dev_evaluator(model)
# 7. Create a trainer & train
trainer = SentenceTransformerTrainer(
model=model,
args=args,
train_dataset=train_dataset,
eval_dataset=eval_dataset,
loss=loss,
evaluator=dev_evaluator,
)
trainer.train()
# (Optional) Evaluate the trained model on the evaluation set once more, this will also log the results
# and include them in the model card
dev_evaluator(model)
queries = dict(enumerate(test_dataset["question"]))
corpus = dict(enumerate(test_dataset["passage_text"] + train_dataset["passage_text"][:30_000]))
relevant_docs = {idx: [idx] for idx in queries}
test_evaluator = InformationRetrievalEvaluator(
queries=queries,
corpus=corpus,
relevant_docs=relevant_docs,
name="miriad-test-1kq-31kd", # 1k questions, 31k passages
show_progress_bar=True,
)
test_evaluator(model)
# 8. Save the trained model
final_output_dir = f"models/{run_name}/final"
model.save_pretrained(final_output_dir)
# 9. (Optional) Push it to the Hugging Face Hub
# It is recommended to run `huggingface-cli login` to log into your Hugging Face account first
try:
model.push_to_hub(run_name)
except Exception:
logging.error(
f"Error uploading model to the Hugging Face Hub:\n{traceback.format_exc()}To upload it manually, you can run "
f"`huggingface-cli login`, followed by loading the model using `model = SentenceTransformer({final_output_dir!r})` "
f"and saving it using `model.push_to_hub('{run_name}')`."
)
训练
我们在配备 24GB 显存的 RTX 3090 上运行了完整的训练脚本,完整的训练和评估脚本共耗时 5.5 小时。如有需要,你可以通过减少 CachedMultipleNegativesRankingLoss 实例上的 mini_batch_size 以及 InformationRetrievalEvaluator 实例上的 batch_size 来进一步降低内存占用。请在此查看我们训练运行的日志:
| Epoch | Step | 训练损失 | 验证损失 | miriad-eval-1kq-31kd_cosine_ndcg@10 | miriad-test-1kq-31kd_cosine_ndcg@10 |
|---|---|---|---|---|---|
| -1 | -1 | - | - | 0.8474 | 0.8340 |
| 0.0256 | 20 | 0.1019 | - | - | - |
| 0.0512 | 40 | 0.0444 | - | - | - |
| 0.0767 | 60 | 0.0408 | - | - | - |
| 0.1023 | 80 | 0.0462 | - | - | - |
| 0.1279 | 100 | 0.0542 | 0.0525 | 0.8616 | - |
| 0.1535 | 120 | 0.0454 | - | - | - |
| 0.1790 | 140 | 0.0403 | - | - | - |
| 0.2046 | 160 | 0.0463 | - | - | - |
| 0.2302 | 180 | 0.0508 | - | - | - |
| 0.2558 | 200 | 0.0497 | 0.0449 | 0.8643 | - |
| 0.2813 | 220 | 0.0451 | - | - | - |
| 0.3069 | 240 | 0.0445 | - | - | - |
| 0.3325 | 260 | 0.0489 | - | - | - |
| 0.3581 | 280 | 0.0452 | - | - | - |
| 0.3836 | 300 | 0.0461 | 0.0406 | 0.8832 | - |
| 0.4092 | 320 | 0.0415 | - | - | - |
| 0.4348 | 340 | 0.04 | - | - | - |
| 0.4604 | 360 | 0.0399 | - | - | - |
| 0.4859 | 380 | 0.0423 | - | - | - |
| 0.5115 | 400 | 0.0352 | 0.0316 | 0.8823 | - |
| 0.5371 | 420 | 0.0408 | - | - | - |
| 0.5627 | 440 | 0.0356 | - | - | - |
| 0.5882 | 460 | 0.0371 | - | - | - |
| 0.6138 | 480 | 0.0276 | - | - | - |
| 0.6394 | 500 | 0.028 | 0.0280 | 0.8807 | - |
| 0.6650 | 520 | 0.0302 | - | - | - |
| 0.6905 | 540 | 0.0345 | - | - | - |
| 0.7161 | 560 | 0.0325 | - | - | - |
| 0.7417 | 580 | 0.033 | - | - | - |
| 0.7673 | 600 | 0.0314 | 0.0264 | 0.8910 | - |
| 0.7928 | 620 | 0.033 | - | - | - |
| 0.8184 | 640 | 0.029 | - | - | - |
| 0.8440 | 660 | 0.0396 | - | - | - |
| 0.8696 | 680 | 0.0266 | - | - | - |
| 0.8951 | 700 | 0.0262 | 0.0240 | 0.8968 | - |
| 0.9207 | 720 | 0.0262 | - | - | - |
| 0.9463 | 740 | 0.0327 | - | - | - |
| 0.9719 | 760 | 0.0293 | - | - | - |
| 0.9974 | 780 | 0.0304 | - | - | - |
| -1 | -1 | - | - | 0.9026 | 0.8862 |
微调评估
基础模型的性能已经非常出色,在我们的 MIRIAD 测试集上取得了 0.8340 的强劲 NDCG@10。尽管如此,我们仍能在这个特定领域数据集上大幅提升它。
| 模型 | 参数量 | miriad-test-1kq-31kd 上的 NDCG@10 |
|---|---|---|
BAAI/bge-base-en-v1.5 |
109M | 0.7541 |
intfloat/multilingual-e5-small |
118M | 0.6852 |
ibm-granite/granite-embedding-125m-english |
125M | 0.7745 |
Snowflake/snowflake-arctic-embed-m-long |
137M | 0.7514 |
intfloat/multilingual-e5-base |
278M | 0.7052 |
Snowflake/snowflake-arctic-embed-m-v2.0 |
305M | 0.8467 |
BAAI/bge-large-en-v1.5 |
335M | 0.7727 |
mixedbread-ai/mxbai-embed-large-v1 |
335M | 0.7851 |
intfloat/multilingual-e5-large |
560M | 0.7318 |
Snowflake/snowflake-arctic-embed-l-v2.0 |
568M | 0.8433 |
Qwen/Qwen3-Embedding-0.6B |
596M | 0.8493 |
google/embeddinggemma-300m(基础版) |
268M | 0.8340 |
sentence-transformers/embeddinggemma-300m-medical(微调版) |
268M | 0.8862 |
我们的微调过程在测试集上实现了 +0.0522 NDCG@10 的显著提升,由此得到的模型在此模型规模下,在我们特定任务上轻松超越了任何现有的通用嵌入模型。额外投入时间和算力将能带来更强的结果,例如 难负例挖掘 或使用超过 10 万对数据进行训练。
延伸阅读
- google/embeddinggemma-300m
- Google EmbeddingGemma 博客文章
- Google EmbeddingGemma 技术报告
- Sentence Transformers 文档
- Sentence Transformers > 训练概览文档
- Transformers.js 文档
- Text Embeddings Inference (TEI) 文档
配套的 Sentence Transformers 博客文章
要扩展本文中的微调方法,或将相关技术应用于 EmbeddingGemma 及其他嵌入模型:
- 使用 Sentence Transformers 训练和微调嵌入模型:本文微调部分所基于的端到端训练指南。
- 使用 Sentence Transformers 训练和微调重排序模型:Cross Encoder 训练,通常用作 EmbeddingGemma 检索器之上的第二阶段。
- 使用 Sentence Transformers 训练和微调稀疏嵌入模型:SPLADE 训练,可用于与 EmbeddingGemma 配合的混合检索。
- 🪆 Matryoshka 嵌入模型简介:EmbeddingGemma 可截断嵌入的背景知识。
- 用于显著更快、更便宜检索的二进制和标量嵌入量化:在检索时进一步缩小 EmbeddingGemma 嵌入。
来源:Hugging Face:Blog · huggingface.co