Как настроить EmbeddingGemma

Посмотреть на ai.google.dev Запустить в Google Colab Запустить в Kaggle Открыть в Vertex AI Посмотреть исходный код на GitHub

Точная настройка помогает устранить разрыв между пониманием модели общего назначения и специализированной высокой точностью, необходимой для вашего приложения. Поскольку ни одна модель не подходит для всех задач, тонкая настройка позволяет адаптировать ее к вашему домену.

Предположим, что ваша компания Shibuya Financial предлагает различные сложные финансовые продукты, такие как инвестиционные фонды, счета NISA (с налоговыми льготами) и ипотечные кредиты. Служба поддержки клиентов использует внутреннюю базу знаний, чтобы быстро находить ответы на вопросы клиентов.

Настройка

Прежде чем начать работу с этим руководством, выполните следующие действия:

  • Чтобы получить доступ к EmbeddingGemma, войдите в аккаунт Hugging Face и выберите Acknowledge license (Принять лицензию) для модели Gemma.
  • Сгенерируйте токен доступа Hugging Face и используйте его для входа в Colab.

Этот блокнот будет выполняться на центральном или графическом процессоре.

Как установить пакеты Python

Установите библиотеки, необходимые для запуска модели EmbeddingGemma и создания встраиваний. Sentence Transformers – это фреймворк на языке Python для встраивания текста и изображений. Дополнительную информацию можно найти в документации по Sentence Transformers.

pip install -U sentence-transformers git+https://github.com/huggingface/transformers@v4.56.0-Embedding-Gemma-preview

После принятия лицензии вам понадобится действительный токен Hugging Face, чтобы получить доступ к модели.

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

Загрузить модель

Используйте библиотеки sentence-transformers, чтобы создать экземпляр класса модели с помощью EmbeddingGemma.

import torch
from sentence_transformers import SentenceTransformer

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()]))

Device: cuda:0
SentenceTransformer(
  (0): Transformer({'max_seq_length': 2048, '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()
)
Total number of parameters in the model: 307581696

Как подготовить набор данных для точной настройки

Это самый важный шаг. Вам нужно создать набор данных, который научит модель понимать, что означает "похожий" в вашем контексте. Эти данные часто структурированы в виде троек: (якорь, положительный пример, отрицательный пример).

  • Якорь – исходный запрос или предложение.
  • Положительный пример – предложение, которое семантически очень похоже на исходное или идентично ему.
  • Отрицательный пример – предложение, которое относится к той же теме, но отличается по смыслу.

В этом примере мы подготовили только три триплета, но для реального приложения вам понадобится гораздо больший набор данных.

from datasets import Dataset

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)

Dataset({
    features: ['anchor', 'positive', 'negative'],
    num_rows: 3
})

До тонкой настройки

По запросу "инвестиции без налога" могут быть получены следующие результаты с оценками сходства:

  1. Document: Opening a NISA account (Документ: открытие счета NISA) (Оценка: 0,51)
  2. Документ: открытие обычного сберегательного счета (оценка: 0,50) <- Похожая оценка, может ввести в заблуждение
  3. Документ: руководство по подаче заявки на ипотечный кредит (оценка: 0,44)
task_name = "STS"

def get_scores(query, documents):
  # Calculate embeddings by calling model.encode()
  query_embeddings = model.encode(query, prompt_name=task_name)
  doc_embeddings = model.encode(documents, prompt_name=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)

Document:  Opening a NISA Account -> 🤖 Score:  0.51571906
Document:  Opening a Regular Savings Account -> 🤖 Score:  0.5035889
Document:  Home Loan Application Guide -> 🤖 Score:  0.4406476

Обучение

Используя фреймворк, например sentence-transformers в Python, базовая модель постепенно изучает тонкие различия в вашей финансовой лексике.

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()

Step 3 finished. Running evaluation:
Document:  Opening a NISA Account -> 🤖 Score:  0.6459116
Document:  Opening a Regular Savings Account -> 🤖 Score:  0.42690125
Document:  Home Loan Application Guide -> 🤖 Score:  0.40419024
Step 6 finished. Running evaluation:
Document:  Opening a NISA Account -> 🤖 Score:  0.68530923
Document:  Opening a Regular Savings Account -> 🤖 Score:  0.3611964
Document:  Home Loan Application Guide -> 🤖 Score:  0.40812016
Step 9 finished. Running evaluation:
Document:  Opening a NISA Account -> 🤖 Score:  0.7168733
Document:  Opening a Regular Savings Account -> 🤖 Score:  0.3449782
Document:  Home Loan Application Guide -> 🤖 Score:  0.44477722
Step 12 finished. Running evaluation:
Document:  Opening a NISA Account -> 🤖 Score:  0.73008573
Document:  Opening a Regular Savings Account -> 🤖 Score:  0.34124148
Document:  Home Loan Application Guide -> 🤖 Score:  0.4676212
Step 15 finished. Running evaluation:
Document:  Opening a NISA Account -> 🤖 Score:  0.73378766
Document:  Opening a Regular Savings Account -> 🤖 Score:  0.34055778
Document:  Home Loan Application Guide -> 🤖 Score:  0.47503752
Step 15 finished. Running evaluation:
Document:  Opening a NISA Account -> 🤖 Score:  0.73378766
Document:  Opening a Regular Savings Account -> 🤖 Score:  0.34055778
Document:  Home Loan Application Guide -> 🤖 Score:  0.47503752
TrainOutput(global_step=15, training_loss=0.009651267528511198, metrics={'train_runtime': 195.3004, 'train_samples_per_second': 0.077, 'train_steps_per_second': 0.077, 'total_flos': 0.0, 'train_loss': 0.009651267528511198, 'epoch': 5.0})

После тонкой настройки

Теперь результаты поиска выглядят гораздо понятнее:

  1. Документ: открытие счета NISA (оценка: 0,73) <- Гораздо более уверенно
  2. Документ: "Открытие обычного сберегательного счета" (оценка: 0,34) <- Явно менее релевантный
  3. Document: Home Loan Application Guide (Score: 0.47)
get_scores(query, documents)

Document:  Opening a NISA Account -> 🤖 Score:  0.73378766
Document:  Opening a Regular Savings Account -> 🤖 Score:  0.34055778
Document:  Home Loan Application Guide -> 🤖 Score:  0.47503752

Чтобы загрузить модель в Hugging Face Hub, можно использовать метод push_to_hub из библиотеки Sentence Transformers.

Загрузив модель, вы сможете легко получать к ней доступ для логического вывода прямо из Хаба, делиться ею с другими и отслеживать версии. После загрузки любой пользователь сможет загрузить вашу модель с помощью одной строки кода, просто указав ее уникальный идентификатор <username>/my-embedding-gemma.

# Push to Hub
model.push_to_hub("my-embedding-gemma")

Выводы и дальнейшие действия

Теперь вы знаете, как адаптировать модель EmbeddingGemma для определенного домена, выполнив ее тонкую настройку с помощью библиотеки Sentence Transformers.

Узнайте, что ещё можно делать с EmbeddingGemma: