Google 发布 EmbeddingGemma 多语言嵌入模型
Welcome EmbeddingGemma, Google's new efficient embedding model
Google 发布 EmbeddingGemma,一个面向端侧场景的多语言嵌入模型,308M 参数、2K 上下文窗口,支持 100 多种语言。该模型基于 Gemma3 骨干并改为双向注意力,量化后内存占用低于 200 MB,作者称其在 500M 参数以下的纯文本多语言嵌入模型中 MTEB 排名最高。
读者可以据此了解 308M 参数多语言嵌入模型在 MTEB 上的位置以及端侧 RAG 的落地路径。
TL;DR
今天,Google 发布 EmbeddingGemma,这是一款最先进的多语言嵌入模型,非常适合端侧使用场景。该模型专为速度和效率而设计,具有 308M 参数的紧凑体积和 2K 上下文窗口,为移动端 RAG 流水线、智能体等解锁了新的可能性。EmbeddingGemma 经过训练可支持超过 100 种语言,并且在撰写本文时,在 Massive Text Embedding Benchmark (MTEB) 上是 500M 以下排名最高的纯文本多语言嵌入模型。
目录
简介
文本嵌入已成为现代自然语言应用的支柱,将单词、句子和文档转化为能够捕捉含义、情感和意图的稠密向量。这些向量支持跨海量语料的快速相似性搜索、聚类、分类和检索,为从推荐引擎和语义搜索到检索增强生成和代码搜索工具的一切提供支持。计算这些嵌入的嵌入模型被广泛使用,在 HF Mirror 上每月下载量远超 2 亿次。
在此基础上,Google DeepMind 的 EmbeddingGemma 作为迄今为止最新、最强大的小型多语言嵌入模型问世。EmbeddingGemma 仅具有 308M 参数、2k token 上下文窗口,并支持超过 100 种语言,在 Massive Multilingual Text Embedding Benchmark (MMTEB) 上实现了最先进的性能,同时在量化后 RAM 占用保持在 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 用法。
该模型使用精心筛选的多语言语料库进行训练,总计约 3200 亿个 token。这一专有数据集融合了公开可用的网络文本、代码和技术文档,以及合成的特定任务示例。它已经过过滤,以避免儿童性虐待材料 (CSAM)、敏感数据以及低质量或不安全内容。
评估
EmbeddingGemma 在 MMTEB(多语言,v2)和 MTEB(英语,v2)测试套件上进行了基准测试,这些套件涵盖了广泛的任务、领域和语言。尽管其参数规模仅为 308M,该模型仍持续超越同类基线,同时保持极小的内存占用。
| MTEB(多语言,v2)性能 | MTEB(英语,v2)性能 |
|---|---|
|
|
结果将列在官方 MTEB 排行榜上。我们排除了任何在超过 20% 的 MTEB 数据上训练过的模型,以减轻潜在的过拟合。
演示
该 演示 还可以 以全屏方式体验。
在桌面设备上亲自体验 演示。
用法
EmbeddingGemma 已与许多流行工具集成,使其易于融入你现有的工作流和应用中。该模型已集成到 Sentence Transformers 中,因此也集成到在后台使用 Sentence Transformers 的项目中,例如 LangChain、LlamaIndex、Haystack 和 txtai。请参阅以下示例,以你偏好的框架开始使用。
对于生产部署,你可以使用 Text Embeddings Inference (TEI) 在各种硬件配置上高效地提供模型服务,并且你可以使用 Transformers.js 用于 Web 应用。
无论你选择哪种框架,都应注意 提示词。对于嵌入模型,提示词会添加到输入文本前面,以便模型区分不同的任务。EmbeddingGemma 是使用这些提示词名称和提示词训练的,因此在使用该模型时也应包含它们:
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/[email protected]
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/[email protected]
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://hf-mirror.com/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/[email protected]
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://hf-mirror.com/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/[email protected]
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/[email protected]
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://hf-mirror.com/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()
# )
这段代码从 HF Mirror 加载 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 函数从 HF Mirror 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 batch sampler 以避免意外的假阴性。
训练参数
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 HF Mirror Hub
# It is recommended to run `huggingface-cli login` to log into your HF Mirror account first
try:
model.push_to_hub(run_name)
except Exception:
logging.error(
f"Error uploading model to the HF Mirror 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 · hf-mirror.com