Gerar embeddings com Sentence Transformers

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

O EmbeddingGemma é um modelo de incorporação leve e aberto projetado para recuperação rápida e de alta qualidade em dispositivos cotidianos, como smartphones. Com apenas 308 milhões de parâmetros, ele é eficiente o suficiente para executar técnicas avançadas de IA, como a geração aumentada por recuperação (RAG, na sigla em inglês), diretamente na sua máquina local sem precisar de uma conexão de Internet.

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

Gerando embedding

Um embedding é uma representação numérica de texto, como uma palavra ou frase, que captura o significado semântico. Basicamente, é uma lista de números (um vetor) que permite que os computadores entendam as relações e o contexto das palavras.

Vamos ver como o EmbeddingGemma processaria três palavras diferentes ["apple", "banana", "car"].

A EmbeddingGemma foi treinada com grandes quantidades de texto e aprendeu as relações entre palavras e conceitos.

words = ["apple", "banana", "car"]

# Calculate embeddings by calling model.encode()
embeddings = model.encode(words)

print(embeddings)
for idx, embedding in enumerate(embeddings):
  print(f"Embedding {idx+1} (shape): {embedding.shape}")

[[-0.18476306  0.00167681  0.03773484 ... -0.07996225 -0.02348064
   0.00976741]
 [-0.21189538 -0.02657359  0.02513712 ... -0.08042689 -0.01999852
   0.00512146]
 [-0.18924113 -0.02551468  0.04486253 ... -0.06377774 -0.03699806
   0.03973572]]
Embedding 1: (768,)
Embedding 2: (768,)
Embedding 3: (768,)

O modelo gera um vetor numérico para cada frase. Os vetores reais são muito longos (768), mas, para simplificar, eles são apresentados com algumas dimensões.

A chave não são os números individuais, mas a distância entre os vetores. Se representarmos esses vetores em um espaço multidimensional, os vetores de apple e banana vão ficar muito próximos um do outro. e o vetor de car ficaria longe dos outros dois.

Como determinar a similaridade

Nesta seção, usamos embeddings para determinar o grau de semelhança semântica entre frases diferentes. Aqui mostramos exemplos com pontuações de similaridade alta, média e baixa.

  • Alta similaridade:

    • Frase A: "O chef preparou uma refeição deliciosa para os convidados".
    • Frase B: "Um jantar delicioso foi preparado pelo chef para os visitantes".
    • Motivo: as duas frases descrevem o mesmo evento usando palavras e estruturas gramaticais diferentes (voz ativa x voz passiva). Elas transmitem o mesmo significado principal.

  • Similaridade média:

    • Frase A: "Ela é especialista em aprendizado de máquina".
    • Frase B: "Ele tem um grande interesse em inteligência artificial."
    • Motivo: as frases estão relacionadas porque o machine learning é uma subárea da inteligência artificial. No entanto, eles falam sobre pessoas diferentes com níveis de engajamento diferentes (especialista x interesse).

  • Baixa semelhança:

    • Frase A: "O tempo em Tóquio está ensolarado hoje".
    • Frase B: "Preciso comprar alimentos para a semana".
    • Motivo: as duas frases são sobre assuntos completamente não relacionados e não compartilham sobreposição semântica.
# The sentences to encode
sentence_high = [
    "The chef prepared a delicious meal for the guests.",
    "A tasty dinner was cooked by the chef for the visitors."
]
sentence_medium = [
    "She is an expert in machine learning.",
    "He has a deep interest in artificial intelligence."
]
sentence_low = [
    "The weather in Tokyo is sunny today.",
    "I need to buy groceries for the week."
]

for sentence in [sentence_high, sentence_medium, sentence_low]:
  print("🙋‍♂️")
  print(sentence)
  embeddings = model.encode(sentence)
  similarities = model.similarity(embeddings[0], embeddings[1])
  print("`-> 🤖 score: ", similarities.numpy()[0][0])

🙋‍♂️
['The chef prepared a delicious meal for the guests.', 'A tasty dinner was cooked by the chef for the visitors.']
`-> 🤖 score:  0.8002148
🙋‍♂️
['She is an expert in machine learning.', 'He has a deep interest in artificial intelligence.']
`-> 🤖 score:  0.45417833
🙋‍♂️
['The weather in Tokyo is sunny today.', 'I need to buy groceries for the week.']
`-> 🤖 score:  0.22262995

Usar comandos com o EmbeddingGemma

Para gerar os melhores embeddings com o EmbeddingGemma, adicione um "comando de instrução" ou uma "tarefa" ao início do texto de entrada. Esses comandos otimizam os embeddings para tarefas específicas, como recuperação de documentos ou resposta a perguntas, e ajudam o modelo a distinguir entre diferentes tipos de entrada, como uma consulta de pesquisa e um documento.

Como usar comandos

Você pode aplicar um comando durante a inferência de três maneiras.

  1. Usando o argumento prompt
    Transmita a string de comando completa diretamente para o método encode. Isso dá a você um controle preciso.

    embeddings = model.encode(
    sentence,
    prompt="task: sentence similarity | query: "
)

  2. Usar o argumento prompt_name
    Selecione um comando predefinido pelo nome. Esses comandos são carregados da configuração do modelo ou durante a inicialização dele.

    embeddings = model.encode(sentence, prompt_name="STS")

  3. Usar o comando padrão
    Se você não especificar prompt nem prompt_name, o sistema vai usar automaticamente o comando definido como default_prompt_name. Se nenhum padrão for definido, nenhum comando será aplicado.

    embeddings = model.encode(sentence)

print("Available tasks:")
for name, prefix in model.prompts.items():
  print(f" {name}: \"{prefix}\"")
print("-"*80)

for sentence in [sentence_high, sentence_medium, sentence_low]:
  print("🙋‍♂️")
  print(sentence)
  embeddings = model.encode(sentence, prompt_name="STS")
  similarities = model.similarity(embeddings[0], embeddings[1])
  print("`-> 🤖 score: ", similarities.numpy()[0][0])

Available tasks:
 query: "task: search result | query: "
 document: "title: none | text: "
 BitextMining: "task: search result | query: "
 Clustering: "task: clustering | query: "
 Classification: "task: classification | query: "
 InstructionRetrieval: "task: code retrieval | query: "
 MultilabelClassification: "task: classification | query: "
 PairClassification: "task: sentence similarity | query: "
 Reranking: "task: search result | query: "
 Retrieval: "task: search result | query: "
 Retrieval-query: "task: search result | query: "
 Retrieval-document: "title: none | text: "
 STS: "task: sentence similarity | query: "
 Summarization: "task: summarization | query: "
--------------------------------------------------------------------------------
🙋‍♂️
['The chef prepared a delicious meal for the guests.', 'A tasty dinner was cooked by the chef for the visitors.']
`-> 🤖 score:  0.9363755
🙋‍♂️
['She is an expert in machine learning.', 'He has a deep interest in artificial intelligence.']
`-> 🤖 score:  0.6425841
🙋‍♂️
['The weather in Tokyo is sunny today.', 'I need to buy groceries for the week.']
`-> 🤖 score:  0.38587403

Caso de uso: geração aumentada por recuperação (RAG)

Para sistemas RAG, use os seguintes valores de prompt_name para criar embeddings especializados para suas consultas e documentos:

  • Para consultas:use prompt_name="Retrieval-query".

    query_embedding = model.encode(
    "How do I use prompts with this model?",
    prompt_name="Retrieval-query"
)

  • Para documentos:use prompt_name="Retrieval-document". Para melhorar ainda mais as incorporações de documentos, inclua um título usando diretamente o argumento prompt:

    • Com um título:
    doc_embedding = model.encode(
    "The document text...",
    prompt="title: Using Prompts in RAG | text: "
)
    • Sem título:
    doc_embedding = model.encode(
    "The document text...",
    prompt="title: none | text: "
)

Leitura adicional

Classificação

A classificação é a tarefa de atribuir um trecho de texto a uma ou mais categorias ou rótulos predefinidos. É uma das tarefas mais fundamentais no processamento de linguagem natural (PLN).

Uma aplicação prática da classificação de texto é o encaminhamento de tíquetes de suporte ao cliente. Esse processo direciona automaticamente as consultas dos clientes para o departamento correto, economizando tempo e reduzindo o trabalho manual.

labels = ["Billing Issue", "Technical Support", "Sales Inquiry"]

sentence = [
  "Excuse me, the app freezes on the login screen. It won't work even when I try to reset my password.",
  "I would like to inquire about your enterprise plan pricing and features for a team of 50 people.",
]

# Calculate embeddings by calling model.encode()
label_embeddings = model.encode(labels, prompt_name="Classification")
embeddings = model.encode(sentence, prompt_name="Classification")

# Calculate the embedding similarities
similarities = model.similarity(embeddings, label_embeddings)
print(similarities)

idx = similarities.argmax(1)
print(idx)

for example in sentence:
  print("🙋‍♂️", example, "-> 🤖", labels[idx[sentence.index(example)]])

tensor([[0.4673, 0.5145, 0.3604],
        [0.4191, 0.5010, 0.5966]])
tensor([1, 2])
🙋‍♂️ Excuse me, the app freezes on the login screen. It won't work even when I try to reset my password. -> 🤖 Technical Support
🙋‍♂️ I would like to inquire about your enterprise plan pricing and features for a team of 50 people. -> 🤖 Sales Inquiry

Aprendizado de representações Matryoshka (MRL)

O EmbeddingGemma usa o MRL para oferecer vários tamanhos de incorporação de um modelo. É um método de treinamento inteligente que cria um único embedding de alta qualidade em que as informações mais importantes se concentram no início do vetor.

Isso significa que você pode ter um embedding menor, mas ainda muito útil, simplesmente usando as primeiras N dimensões do embedding completo. Usar embeddings menores e truncados é muito mais barato para armazenar e mais rápido para processar, mas essa eficiência tem um custo: a qualidade dos embeddings pode ser menor. Com a MRL, você pode escolher o equilíbrio ideal entre velocidade e precisão para as necessidades específicas do seu aplicativo.

Vamos usar três palavras ["apple", "banana", "car"] e criar embeddings simplificados para ver como o MRL funciona.

def check_word_similarities():
  # Calculate the embedding similarities
  print("similarity function: ", model.similarity_fn_name)
  similarities = model.similarity(embeddings[0], embeddings[1:])
  print(similarities)

  for idx, word in enumerate(words[1:]):
    print("🙋‍♂️ apple vs.", word, "-> 🤖 score: ", similarities.numpy()[0][idx])

# Calculate embeddings by calling model.encode()
embeddings = model.encode(words, prompt_name="STS")

check_word_similarities()

similarity function:  cosine
tensor([[0.7510, 0.6685]])
🙋‍♂️ apple vs. banana -> 🤖 score:  0.75102395
🙋‍♂️ apple vs. car -> 🤖 score:  0.6684626

Agora, para uma aplicação mais rápida, não é necessário um novo modelo. Basta truncar os embeddings completos nas primeiras 512 dimensões. Para resultados ideais, também recomendamos definir normalize_embeddings=True, que dimensiona os vetores para um comprimento de unidade de 1.

embeddings = model.encode(words, truncate_dim=512, normalize_embeddings=True)

for idx, embedding in enumerate(embeddings):
  print(f"Embedding {idx+1}: {embedding.shape}")

print("-"*80)
check_word_similarities()

Embedding 1: (512,)
Embedding 2: (512,)
Embedding 3: (512,)
--------------------------------------------------------------------------------
similarity function:  cosine
tensor([[0.7674, 0.7041]])
🙋‍♂️ apple vs. banana -> 🤖 score:  0.767427
🙋‍♂️ apple vs. car -> 🤖 score:  0.7040509

Em ambientes extremamente restritos, é possível reduzir ainda mais as incorporações para apenas 256 dimensões. Você também pode usar o produto escalar, que é mais eficiente, para cálculos de similaridade em vez da similaridade de cosseno padrão.

model = SentenceTransformer(model_id, truncate_dim=256, similarity_fn_name="dot").to(device=device)
embeddings = model.encode(words, prompt_name="STS", normalize_embeddings=True)

for idx, embedding in enumerate(embeddings):
  print(f"Embedding {idx+1}: {embedding.shape}")

print("-"*80)
check_word_similarities()

Embedding 1: (256,)
Embedding 2: (256,)
Embedding 3: (256,)
--------------------------------------------------------------------------------
similarity function:  dot
tensor([[0.7855, 0.7382]])
🙋‍♂️ apple vs. banana -> 🤖 score:  0.7854644
🙋‍♂️ apple vs. car -> 🤖 score:  0.7382126

Resumo e próximas etapas

Agora você pode gerar embeddings de texto de alta qualidade usando o EmbeddingGemma e a biblioteca Sentence Transformers. Aplique essas habilidades para criar recursos avançados, como similaridade semântica, classificação de texto e sistemas de geração aumentada de recuperação (RAG), e continue explorando o que é possível fazer com os modelos da Gemma.

Confira os seguintes documentos: