Тонкая настройка EmbeddingGemma

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

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

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

Настраивать

Перед началом работы с этим руководством выполните следующие шаги:

  • Чтобы получить доступ к EmbeddingGemma, войдите в Hugging Face и выберите «Подтвердить лицензию» для модели 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

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

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

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

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

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. Документ: Открытие счета 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. Документ: Руководство по подаче заявления на получение ипотечного кредита (оценка: 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.

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

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

Резюме и дальнейшие шаги

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

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