1
# 示例：不同维度的嵌入使用
2
from sentence_transformers import SentenceTransformer
3

4
model = SentenceTransformer("google/embeddinggemma-300m")
5

6
# 生成完整768维嵌入
7
full_embedding = model.encode(["示例文本"], output_dimensions=768)
8

9
# 截断为256维嵌入，适合快速检索
10
compact_embedding = model.encode(["示例文本"], output_dimensions=256)

1
# 本地RAG示例
2
import numpy as np
3
from sentence_transformers import SentenceTransformer
4
from sklearn.metrics.pairwise import cosine_similarity
5

6
# 加载EmbeddingGemma模型
7
embedding_model = SentenceTransformer("google/embeddinggemma-300m")
8

9
# 文档库嵌入
10
documents = [
11
    "人工智能是计算机科学的一个分支",
12
    "机器学习是实现人工智能的重要方法",
13
    "深度学习是机器学习的一个子集"
14
]
15

16
doc_embeddings = embedding_model.encode(documents)
17

18
# 查询处理
19
query = "什么是AI？"
20
query_embedding = embedding_model.encode([query])
21

22
# 相似度搜索
23
similarities = cosine_similarity(query_embedding, doc_embeddings)
24
best_match_idx = np.argmax(similarities)
25

26
print(f"最相关文档: {documents[best_match_idx]}")

1
# 安装必要的依赖
2
pip install sentence-transformers
3
pip install transformers
4
pip install torch

1
from sentence_transformers import SentenceTransformer
2

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

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

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

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

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

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

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

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

1
# 与LangChain集成
2
from langchain.embeddings import HuggingFaceEmbeddings
3

4
embeddings = HuggingFaceEmbeddings(
5
    model_name="google/embeddinggemma-300m"
6
)
7

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

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

1
# 加载模型
2
import torch
3
from sentence_transformers import SentenceTransformer
4
from datasets import Dataset
5

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

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

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

15
# 准备微调数据集
16
dataset = [
17
    ["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."],
18
    ["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?"],
19
    ["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?"],
20
]
21

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

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

29
task_name = "STS"
30

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

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

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

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

45
get_scores(query, documents)
46

47
from sentence_transformers import SentenceTransformerTrainer, SentenceTransformerTrainingArguments
48
from sentence_transformers.losses import MultipleNegativesRankingLoss
49
from transformers import TrainerCallback
50

51
loss = MultipleNegativesRankingLoss(model)
52

53
args = SentenceTransformerTrainingArguments(
54
    # Required parameter:
55
    output_dir="my-embedding-gemma",
56
    # Optional training parameters:
57
    prompts=model.prompts[task_name],    # use model's prompt to train
58
    num_train_epochs=5,
59
    per_device_train_batch_size=1,
60
    learning_rate=2e-5,
61
    warmup_ratio=0.1,
62
    # Optional tracking/debugging parameters:
63
    logging_steps=train_dataset.num_rows,
64
    report_to="none",
65
)
66

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

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

77
def evaluate():
78
  get_scores(query, documents)
79

80
trainer = SentenceTransformerTrainer(
81
    model=model,
82
    args=args,
83
    train_dataset=train_dataset,
84
    loss=loss,
85
    callbacks=[MyCallback(evaluate)]
86
)
87
trainer.train()

1
# CPU部署
2
docker run -p 8080:80 ghcr.io/huggingface/text-embeddings-inference:cpu-1.8.1 \
3
  --model-id google/embeddinggemma-300m --dtype float32
4

5
# GPU部署（支持多种GPU架构）
6
docker run --gpus all --shm-size 1g -p 8080:80 \
7
  ghcr.io/huggingface/text-embeddings-inference:cuda-1.8.1 \
8
  --model-id google/embeddinggemma-300m --dtype float32

1
# 使用ONNX优化版本
2
docker run -p 8080:80 ghcr.io/huggingface/text-embeddings-inference:cpu-1.8.1 \
3
  --model-id onnx-community/embeddinggemma-300m-ONNX \
4
  --dtype float32 --pooling mean

1
from sentence_transformers import SentenceTransformer, InputExample, losses
2
from torch.utils.data import DataLoader
3

4
# 准备训练数据
5
train_examples = [
6
    InputExample(texts=['查询文本', '相关文档'], label=1.0),
7
    InputExample(texts=['查询文本', '不相关文档'], label=0.0),
8
]
9

10
# 创建数据加载器
11
train_dataloader = DataLoader(train_examples, shuffle=True, batch_size=16)
12

13
# 加载预训练模型
14
model = SentenceTransformer('google/embeddinggemma-300m')
15

16
# 定义损失函数
17
train_loss = losses.CosineSimilarityLoss(model)
18

19
# 微调模型
20
model.fit(
21
    train_objectives=[(train_dataloader, train_loss)],
22
    epochs=1,
23
    warmup_steps=100
24
)