Genera embeddings con Sentence Transformers

Ver en ai.google.dev Ejecutar en Google Colab Ejecutar en Kaggle Abrir en Vertex AI Ver código fuente en GitHub

EmbeddingGemma es un modelo de incorporación ligero y abierto diseñado para la recuperación rápida y de alta calidad en dispositivos cotidianos, como teléfonos celulares. Con solo 308 millones de parámetros, es lo suficientemente eficiente como para ejecutar técnicas avanzadas de IA, como la generación aumentada por recuperación (RAG), directamente en tu máquina local sin necesidad de conexión a Internet.

Configuración

Antes de comenzar este instructivo, completa los siguientes pasos:

  • Para acceder a Gemma, ingresa a Hugging Face y selecciona Acknowledge license para un modelo de Gemma.
  • Genera un token de acceso de Hugging Face y úsalo para acceder desde Colab.

Este notebook se ejecutará en la CPU o la GPU.

Instala paquetes de Python

Instala las bibliotecas necesarias para ejecutar el modelo EmbeddingGemma y generar embeddings. Sentence Transformers es un framework de Python para incorporaciones de texto e imágenes. Para obtener más información, consulta la documentación de Sentence Transformers.

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

Después de aceptar la licencia, necesitarás un token de Hugging Face válido para acceder al modelo.

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

Cargar modelo

Usa las bibliotecas de sentence-transformers para crear una instancia de una clase de modelo con 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

Generación de embeddings

Una incorporación es una representación numérica de texto, como una palabra o una oración, que captura su significado semántico. Básicamente, es una lista de números (un vector) que permite que las computadoras comprendan las relaciones y el contexto de las palabras.

Veamos cómo EmbeddingGemma procesaría tres palabras diferentes ["apple", "banana", "car"].

EmbeddingGemma se entrenó con grandes cantidades de texto y aprendió las relaciones entre palabras y conceptos.

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

El modelo genera un vector numérico para cada oración. Los vectores reales son muy largos (768), pero, para simplificar, se presentan con algunas dimensiones.

La clave no son los números individuales en sí, sino la distancia entre los vectores. Si graficáramos estos vectores en un espacio multidimensional, los vectores de apple y banana estarían muy cerca entre sí. Y el vector de car estaría lejos de los otros dos.

Cómo determinar la similitud

En esta sección, usaremos las incorporaciones para determinar qué tan similares son semánticamente las diferentes oraciones. Aquí mostramos ejemplos con puntuaciones de similitud altas, medias y bajas.

  • Alta similitud:

    • Oración A: "El chef preparó una comida deliciosa para los invitados".
    • Oración B: "El chef preparó una cena deliciosa para los visitantes".
    • Explicación: Ambas oraciones describen el mismo evento con diferentes palabras y estructuras gramaticales (voz activa y voz pasiva). Transmiten el mismo significado principal.
  • Similitud media:

    • Oración A: "Ella es experta en aprendizaje automático".
    • Oración B: "Tiene un gran interés en la inteligencia artificial".
    • Justificación: Las oraciones están relacionadas, ya que el aprendizaje automático es un subcampo de la inteligencia artificial. Sin embargo, hablan de personas diferentes con distintos niveles de participación (expertos vs. interesados).
  • Similitud baja:

    • Oración A: "El clima en Tokio es soleado hoy".
    • Oración B: "Necesito comprar alimentos para la semana".
    • Explicación: Las dos oraciones tratan temas completamente no relacionados y no comparten ninguna superposición 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

Cómo usar instrucciones con EmbeddingGemma

Para generar las mejores incorporaciones con EmbeddingGemma, debes agregar una "instrucción" o "tarea" al comienzo del texto de entrada. Estas instrucciones optimizan los embeddings para tareas específicas, como la recuperación de documentos o la búsqueda de respuestas, y ayudan al modelo a distinguir entre diferentes tipos de entrada, como una búsqueda y un documento.

Cómo aplicar instrucciones

Puedes aplicar una instrucción durante la inferencia de tres maneras.

  1. Usa el argumento prompt.
    Pasa la cadena de instrucciones completa directamente al método encode. Esto te brinda un control preciso.

    embeddings = model.encode(
        sentence,
        prompt="task: sentence similarity | query: "
    )
    
  2. Usa el argumento prompt_name
    Selecciona una instrucción predefinida por su nombre. Estas instrucciones se cargan desde la configuración del modelo o durante su inicialización.

    embeddings = model.encode(sentence, prompt_name="STS")
    
  3. Cómo usar la instrucción predeterminada
    Si no especificas prompt ni prompt_name, el sistema usará automáticamente la instrucción establecida como default_prompt_name. Si no se establece una instrucción predeterminada, no se aplicará ninguna instrucción.

    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: Generación mejorada por recuperación (RAG)

En el caso de los sistemas RAG, usa los siguientes valores de prompt_name para crear incorporaciones especializadas para tus búsquedas y documentos:

  • Para las búsquedas: Usa prompt_name="Retrieval-query".

    query_embedding = model.encode(
        "How do I use prompts with this model?",
        prompt_name="Retrieval-query"
    )
    
  • Para documentos: Usa prompt_name="Retrieval-document". Para mejorar aún más las incorporaciones de documentos, también puedes incluir un título usando el argumento prompt directamente:

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

Lecturas adicionales

Clasificación

La clasificación es la tarea de asignar un fragmento de texto a una o más categorías o etiquetas predefinidas. Es una de las tareas más fundamentales del procesamiento de lenguaje natural (PNL).

Una aplicación práctica de la clasificación de texto es el enrutamiento de tickets de asistencia al cliente. Este proceso dirige automáticamente las consultas de los clientes al departamento correcto, lo que ahorra tiempo y reduce el trabajo 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

Aprendizaje de representaciones de Matryoshka (MRL)

EmbeddingGemma aprovecha MRL para proporcionar varios tamaños de incorporación desde un solo modelo. Es un método de entrenamiento inteligente que crea un embedding único y de alta calidad en el que la información más importante se concentra al principio del vector.

Esto significa que puedes obtener una incorporación más pequeña, pero igual de útil, simplemente tomando las primeras N dimensiones de la incorporación completa. El uso de incorporaciones más pequeñas y truncadas es significativamente más económico de almacenar y más rápido de procesar, pero esta eficiencia tiene el costo de una posible menor calidad de las incorporaciones. MRL te permite elegir el equilibrio óptimo entre esta velocidad y la precisión para las necesidades específicas de tu aplicación.

Usemos tres palabras ["apple", "banana", "car"] y creemos incorporaciones simplificadas para ver cómo funciona MRL.

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

Ahora, para una aplicación más rápida, no necesitas un modelo nuevo. Simplemente trunca las incorporaciones completas a las primeras 512 dimensiones. Para obtener resultados óptimos, también se recomienda establecer normalize_embeddings=True, que ajusta los vectores a una longitud unitaria 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

En entornos extremadamente restringidos, puedes acortar aún más las incorporaciones a solo 256 dimensiones. También puedes usar el producto escalar más eficiente para los cálculos de similitud en lugar de la similitud de coseno estándar.

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

Resumen y próximos pasos

Ahora tienes las herramientas necesarias para generar incorporaciones de texto de alta calidad con EmbeddingGemma y la biblioteca de Sentence Transformers. Aplica estas habilidades para crear funciones potentes, como la similitud semántica, la clasificación de texto y los sistemas de generación mejorada por recuperación (RAG), y sigue explorando las posibilidades de los modelos de Gemma.

A continuación, consulta los siguientes documentos: