Ajustar o EmbeddingGemma

Ver em ai.google.dev Executar no Google Colab Executar no Kaggle Abrir na Vertex AI Ver código-fonte no GitHub

O ajuste fino ajuda a diminuir a lacuna entre a compreensão de uso geral de um modelo e a precisão especializada e de alto desempenho que seu aplicativo exige. Como nenhum modelo é perfeito para todas as tarefas, o ajuste fino o adapta ao seu domínio específico.

Imagine que sua empresa, "Shibuya Financial", oferece vários produtos financeiros complexos, como fundos de investimento, contas NISA (uma conta de poupança com benefícios fiscais) e empréstimos imobiliários. Sua equipe de suporte ao cliente usa uma base de conhecimento interna para encontrar respostas rapidamente para as perguntas dos clientes.

Configuração

Antes de iniciar este tutorial, conclua as seguintes etapas:

  • Faça login no Hugging Face e selecione Confirmar licença para um modelo do Gemma.
  • Gere um token de acesso do Hugging Face e use-o para fazer login no Colab.

Este notebook será executado na CPU ou na GPU.

Instalar pacotes Python

Instale as bibliotecas necessárias para executar o modelo EmbeddingGemma e gerar embeddings. O Sentence Transformers é um framework Python para embeddings de texto e imagem. Para mais informações, consulte a documentação do Sentence Transformers (em inglês).

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

Depois de aceitar a licença, você vai precisar de um token do Hugging Face válido para acessar o modelo.

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

Carregar modelo

Use as bibliotecas sentence-transformers para criar uma instância de uma classe de modelo com 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

Preparar o conjunto de dados para ajuste refinado

Essa é a parte mais importante. Você precisa criar um conjunto de dados que ensine ao modelo o que "semelhante" significa no seu contexto específico. Esses dados geralmente são estruturados como trios: (âncora, positivo, negativo)

  • Âncora: a consulta ou frase original.
  • Positiva: uma frase semanticamente muito semelhante ou idêntica à âncora.
  • Negativa: uma frase sobre um tema relacionado, mas semanticamente diferente.

Neste exemplo, preparamos apenas três tripletes, mas, para um aplicativo real, você precisaria de um conjunto de dados muito maior para ter um bom desempenho.

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

Antes do ajuste fino

Uma pesquisa por "investimento livre de tributos" pode ter gerado os seguintes resultados, com pontuações de similaridade:

  1. Documento: Opening a NISA account (Pontuação: 0,45)
  2. Documento: Opening a Regular Saving Account (Pontuação: 0,48) <- Pontuação semelhante, potencialmente confusa
  3. Documento: Guia de solicitação de empréstimo imobiliário (pontuação: 0,42)
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)

Document:  Opening a NISA Account -> 🤖 Score:  0.45698774
Document:  Opening a Regular Savings Account -> 🤖 Score:  0.48092696
Document:  Home Loan Application Guide -> 🤖 Score:  0.42127067

Treinamento

Usando uma estrutura como sentence-transformers em Python, o modelo de base aprende gradualmente as distinções sutis no seu vocabulário financeiro.

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.6449194
Document:  Opening a Regular Savings Account -> 🤖 Score:  0.44123
Document:  Home Loan Application Guide -> 🤖 Score:  0.46752414
Step 6 finished. Running evaluation:
Document:  Opening a NISA Account -> 🤖 Score:  0.68873787
Document:  Opening a Regular Savings Account -> 🤖 Score:  0.34069622
Document:  Home Loan Application Guide -> 🤖 Score:  0.50065553
Step 9 finished. Running evaluation:
Document:  Opening a NISA Account -> 🤖 Score:  0.7148906
Document:  Opening a Regular Savings Account -> 🤖 Score:  0.30480516
Document:  Home Loan Application Guide -> 🤖 Score:  0.52454984
Step 12 finished. Running evaluation:
Document:  Opening a NISA Account -> 🤖 Score:  0.72614634
Document:  Opening a Regular Savings Account -> 🤖 Score:  0.29255486
Document:  Home Loan Application Guide -> 🤖 Score:  0.5370023
Step 15 finished. Running evaluation:
Document:  Opening a NISA Account -> 🤖 Score:  0.7294032
Document:  Opening a Regular Savings Account -> 🤖 Score:  0.2893038
Document:  Home Loan Application Guide -> 🤖 Score:  0.54087913
Step 15 finished. Running evaluation:
Document:  Opening a NISA Account -> 🤖 Score:  0.7294032
Document:  Opening a Regular Savings Account -> 🤖 Score:  0.2893038
Document:  Home Loan Application Guide -> 🤖 Score:  0.54087913
TrainOutput(global_step=15, training_loss=0.009651281436261646, metrics={'train_runtime': 63.2486, 'train_samples_per_second': 0.237, 'train_steps_per_second': 0.237, 'total_flos': 0.0, 'train_loss': 0.009651281436261646, 'epoch': 5.0})

Após o ajuste de detalhes

A mesma pesquisa agora gera resultados muito mais claros:

  1. Documento: como abrir uma conta NISA (pontuação: 0,72) <- Muito mais confiante
  2. Documento: Opening a Regular Saving Account (Pontuação: 0,28) <- Claramente menos relevante
  3. Documento: Guia de solicitação de empréstimo imobiliário (pontuação: 0,54)
get_scores(query, documents)

Document:  Opening a NISA Account -> 🤖 Score:  0.7294032
Document:  Opening a Regular Savings Account -> 🤖 Score:  0.2893038
Document:  Home Loan Application Guide -> 🤖 Score:  0.54087913

Para fazer upload do modelo para o Hugging Face Hub, use o método push_to_hub da biblioteca Sentence Transformers.

Ao fazer upload do seu modelo, fica mais fácil acessar a inferência diretamente do Hub, compartilhar com outras pessoas e controlar as versões do seu trabalho. Depois de fazer o upload, qualquer pessoa pode carregar seu modelo com uma única linha de código, basta referenciar o ID exclusivo do modelo <username>/my-embedding-gemma.

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

Resumo e próximas etapas

Agora você aprendeu a adaptar um modelo EmbeddingGemma para um domínio específico ajustando-o com a biblioteca Sentence Transformers.

Confira o que mais você pode fazer com o EmbeddingGemma: