使用 Sentence Transformers 微調 EmbeddingGemma¶

# Login into Hugging Face Hub
from huggingface_hub import login
login()

from sentence_transformers import SentenceTransformer

model_id = "google/embeddinggemma-2"
model = SentenceTransformer(
    model_id,
    config_kwargs={"vision_config": None, "audio_config": None},
)

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

task_name = "STS"

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",
]

def evaluate_retrieval(query_text, docs):
    query_emb = model.encode(query_text, prompt_name=task_name)
    doc_embs = model.encode(docs, prompt_name=task_name)
    similarities = model.similarity(query_emb, doc_embs)[0]

    for doc, score in zip(docs, similarities):
        print(f"Document: {doc:<35} -> 🤖 Score: {score.item():.4f}")

print("Baseline Scores (Before Fine-Tuning):")
evaluate_retrieval(query, documents)

from datasets import Dataset

triplets = [
    [
        "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 into a Hugging Face Dataset with 'anchor', 'positive', and 'negative' columns
data_as_dicts = [
    {"anchor": row[0], "positive": row[1], "negative": row[2]}
    for row in triplets
]
train_dataset = Dataset.from_list(data_as_dicts)
print(train_dataset)

from sentence_transformers import (
    SentenceTransformerTrainer,
    SentenceTransformerTrainingArguments,
)
from sentence_transformers.sentence_transformer.losses import MultipleNegativesRankingLoss
from transformers import TrainerCallback

# Initialize contrastive ranking loss
loss = MultipleNegativesRankingLoss(model)

# Configure training arguments
args = SentenceTransformerTrainingArguments(
    output_dir="my-embeddinggemma-2",
    prompts=model.prompts[task_name],  # Train with the target task prompt prefix
    num_train_epochs=5,
    per_device_train_batch_size=1,
    learning_rate=2e-5,
    logging_steps=train_dataset.num_rows,
    report_to="none",
)

class EvaluationCallback(TrainerCallback):
    """Callback that evaluates retrieval scores at each logging step."""

    def __init__(self, eval_fn):
        self.eval_fn = eval_fn

    def on_log(self, args, state, control, **kwargs):
        print(f"\n--- Step {state.global_step} Evaluation ---")
        self.eval_fn()

trainer = SentenceTransformerTrainer(
    model=model,
    args=args,
    train_dataset=train_dataset,
    loss=loss,
    callbacks=[EvaluationCallback(lambda: evaluate_retrieval(query, documents))],
)

trainer.train()

print("Post-Training Scores (After Fine-Tuning):")
evaluate_retrieval(query, documents)

# Save locally
model.save_pretrained("my-embeddinggemma-2-finetuned")

# Optional: Push your fine-tuned model to the Hugging Face Hub
model.push_to_hub("my-embeddinggemma-2-finetuned")