Google发布最新开放文本嵌入模型：EmbeddingGemma

开发集成指南

环境配置

# 安装必要的依赖
pip install sentence-transformers
pip install transformers
pip install torch

基础使用示例

from sentence_transformers import SentenceTransformer

# 1. 加载模型
model = SentenceTransformer("google/embeddinggemma-300m")

# 2. 生成嵌入
texts = [
    "今天天气真好",
    "我喜欢机器学习",
    "Python是一门强大的编程语言"
]

embeddings = model.encode(texts)
print(f"嵌入维度: {embeddings.shape}")

# 3. 计算文本相似度
from sklearn.metrics.pairwise import cosine_similarity
similarity_matrix = cosine_similarity(embeddings)
print("相似度矩阵:")
print(similarity_matrix)

针对特定任务的提示模板

EmbeddingGemma支持基于任务的提示模板，可以针对不同的使用场景生成优化的嵌入：

# 查询提示模板
def format_query(query, task="search result"):
    return f"task: {task} | query: {query}"

# 文档提示模板  
def format_document(doc, task="search result"):
    return f"task: {task} | document: {doc}"

# 使用示例
query = format_query("人工智能的发展历史")
document = format_document("人工智能技术在过去几十年中取得了巨大进展")

query_emb = model.encode([query])
doc_emb = model.encode([document])

与主流框架集成

EmbeddingGemma已经与多个流行的AI开发框架无缝集成：

# 与LangChain集成
from langchain.embeddings import HuggingFaceEmbeddings

embeddings = HuggingFaceEmbeddings(
    model_name="google/embeddinggemma-300m"
)

# 与LlamaIndex集成
from llama_index.embeddings.huggingface import HuggingFaceEmbedding

embed_model = HuggingFaceEmbedding(
    model_name="google/embeddinggemma-300m"
)

微调（Fine-tuning）

虽然 EmbeddingGemma 的预训练版本已经足够强大，但在特定领域或专业任务上，通过微调可以让其性能更上一层楼。例如，你可以使用自己业务场景中的数据（如内部知识库的问答对）来对模型进行微调，使其更懂你的“行话”。

官方推荐使用 sentence-transformers 库来进行微调，因为它提供了非常便捷的训练流程。

# 加载模型
import torch
from sentence_transformers import SentenceTransformer
from datasets import Dataset

device = "cuda" if torch.cuda.is_available() else "cpu"

model_id = "google/embeddinggemma-300M"
model = SentenceTransformer(model_id).to(device=device)

print(f"Device: {model.device}")
print(model)
print("Total number of parameters in the model:", sum([p.numel() for _, p in model.named_parameters()]))

# 准备微调数据集
dataset = [
    ["How do I open a NISA account?", "What is the procedure for starting a new tax-free investment account?", "I want to check the balance of my regular savings account."],
    ["Are there fees for making an early repayment on a home loan?", "If I pay back my house loan early, will there be any costs?", "What is the management fee for this investment trust?"],
    ["What is the coverage for medical insurance?", "Tell me about the benefits of the health insurance plan.", "What is the cancellation policy for my life insurance?"],
]

# Convert the list-based dataset into a list of dictionaries.
data_as_dicts = [ {"anchor": row[0], "positive": row[1], "negative": row[2]} for row in dataset ]

# Create a Hugging Face `Dataset` object from the list of dictionaries.
train_dataset = Dataset.from_list(data_as_dicts)
print(train_dataset)

task_name = "STS"

def get_scores(query, documents):
  # Calculate embeddings by calling model.encode()
  query_embeddings = model.encode(query, prompt=task_name)
  doc_embeddings = model.encode(documents, prompt=task_name)

  # Calculate the embedding similarities
  similarities = model.similarity(query_embeddings, doc_embeddings)

  for idx, doc in enumerate(documents):
    print("Document: ", doc, "-> 🤖 Score: ", similarities.numpy()[0][idx])

query = "I want to start a tax-free installment investment, what should I do?"
documents = ["Opening a NISA Account", "Opening a Regular Savings Account", "Home Loan Application Guide"]

get_scores(query, documents)

from sentence_transformers import SentenceTransformerTrainer, SentenceTransformerTrainingArguments
from sentence_transformers.losses import MultipleNegativesRankingLoss
from transformers import TrainerCallback

loss = MultipleNegativesRankingLoss(model)

args = SentenceTransformerTrainingArguments(
    # Required parameter:
    output_dir="my-embedding-gemma",
    # Optional training parameters:
    prompts=model.prompts[task_name],    # use model's prompt to train
    num_train_epochs=5,
    per_device_train_batch_size=1,
    learning_rate=2e-5,
    warmup_ratio=0.1,
    # Optional tracking/debugging parameters:
    logging_steps=train_dataset.num_rows,
    report_to="none",
)

class MyCallback(TrainerCallback):
    "A callback that evaluates the model at the end of eopch"
    def __init__(self, evaluate):
        self.evaluate = evaluate # evaluate function

    def on_log(self, args, state, control, **kwargs):
        # Evaluate the model using text generation
        print(f"Step {state.global_step} finished. Running evaluation:")
        self.evaluate()

def evaluate():
  get_scores(query, documents)

trainer = SentenceTransformerTrainer(
    model=model,
    args=args,
    train_dataset=train_dataset,
    loss=loss,
    callbacks=[MyCallback(evaluate)]
)
trainer.train()