Générer des embeddings avec Sentence Transformers

 Afficher sur ai.google.dev Exécuter dans Google Colab  Exécuter dans Kaggle Ouvrir dans Vertex AI  Afficher la source sur GitHub

EmbeddingGemma est un modèle d'embedding léger et ouvert conçu pour une récupération rapide et de haute qualité sur les appareils du quotidien tels que les téléphones mobiles. Avec seulement 308 millions de paramètres, il est suffisamment efficace pour exécuter des techniques d'IA avancées, telles que la génération augmentée par récupération (RAG, Retrieval Augmented Generation), directement sur votre machine locale sans connexion Internet.

Configuration

Avant de commencer ce tutoriel, effectuez les étapes suivantes :

  • Pour accéder à Gemma, connectez-vous à Hugging Face et sélectionnez Accepter la licence pour un modèle Gemma.
  • Générez un jeton d'accès Hugging Face et utilisez-le pour vous connecter depuis Colab.

Ce notebook s'exécutera sur un processeur ou un GPU.

Installer des packages Python

Installez les bibliothèques requises pour exécuter le modèle EmbeddingGemma et générer des embeddings. Sentence Transformers est un framework Python pour les embeddings de texte et d'image. Pour en savoir plus, consultez la documentation Sentence Transformers.

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

Une fois la licence acceptée, vous avez besoin d'un jeton Hugging Face valide pour accéder au modèle.

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

Charger le modèle

Utilisez les bibliothèques sentence-transformers pour créer une instance d'une classe de modèle avec 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

Génération d'embeddings

Un embedding est une représentation numérique d'un texte (mot ou phrase, par exemple) qui capture sa signification sémantique. Il s'agit essentiellement d'une liste de nombres (un vecteur) qui permet aux ordinateurs de comprendre les relations et le contexte des mots.

Voyons comment EmbeddingGemma traiterait trois mots différents ["apple", "banana", "car"].

EmbeddingGemma a été entraîné sur de grandes quantités de texte et a appris les relations entre les mots et les concepts.

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

Le modèle génère un vecteur numérique pour chaque phrase. Les vecteurs réels sont très longs (768), mais pour plus de simplicité, ils sont présentés avec quelques dimensions.

La clé n'est pas les nombres individuels eux-mêmes, mais la distance entre les vecteurs. Si nous devions représenter ces vecteurs dans un espace multidimensionnel, les vecteurs de apple et banana seraient très proches les uns des autres. Le vecteur de car serait éloigné des deux autres.

Déterminer la similarité

Dans cette section, nous utilisons des embeddings pour déterminer la similarité sémantique entre différentes phrases. Nous présentons ici des exemples avec des scores de similarité élevés, moyens et faibles.

  • Similarité élevée :

    • Phrase A : "Le chef a préparé un délicieux repas pour les invités."
    • Phrase B : "Un délicieux dîner a été préparé par le chef pour les visiteurs."
    • Raisonnement : les deux phrases décrivent le même événement en utilisant des mots et des structures grammaticales différents (voix active ou passive). Elles transmettent le même sens fondamental.

  • Similarité moyenne :

    • Phrase A : "Elle est experte en machine learning."
    • Phrase B : "Il s'intéresse beaucoup à l'intelligence artificielle."
    • Raisonnement : les phrases sont liées, car le machine learning est un sous-domaine de l'intelligence artificielle. Toutefois, elles parlent de personnes différentes avec des niveaux d'engagement différents (expert vs intérêt).

  • Similarité faible :

    • Phrase A : "Il fait beau à Tokyo aujourd'hui."
    • Phrase B : "Je dois acheter des produits alimentaires pour la semaine."
    • Raisonnement : Les deux phrases portent sur des sujets complètement différents et ne présentent aucune similitude sémantique.
# 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

Utiliser des prompts avec EmbeddingGemma

Pour générer les meilleurs embeddings avec EmbeddingGemma, vous devez ajouter une "invite d'instruction" ou une "tâche" au début de votre texte d'entrée. Ces requêtes optimisent les embeddings pour des tâches spécifiques, telles que la récupération de documents ou les questions-réponses, et aident le modèle à faire la distinction entre différents types d'entrées, comme une requête de recherche et un document.

Appliquer des requêtes

Vous pouvez appliquer un prompt lors de l'inférence de trois manières différentes.

  1. Utiliser l'argument prompt
    Transmettez la chaîne de requête complète directement à la méthode encode. Cela vous permet de contrôler précisément ce qui est affiché.

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

  2. Utiliser l'argument prompt_name
    Sélectionnez une invite prédéfinie par son nom. Ces requêtes sont chargées à partir de la configuration du modèle ou lors de son initialisation.

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

  3. Utiliser la requête par défaut
    Si vous ne spécifiez pas prompt ni prompt_name, le système utilisera automatiquement la requête définie comme default_prompt_name. Si aucune valeur par défaut n'est définie, aucune requête n'est appliquée.

    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

Cas d'utilisation : génération augmentée par récupération (RAG)

Pour les systèmes RAG, utilisez les valeurs prompt_name suivantes pour créer des embeddings spécialisés pour vos requêtes et vos documents :

  • Pour les requêtes : utilisez prompt_name="Retrieval-query".

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

  • Pour les documents : utilisez prompt_name="Retrieval-document". Pour améliorer encore les embeddings de documents, vous pouvez également inclure un titre en utilisant directement l'argument prompt :

    • Avec un titre :
    doc_embedding = model.encode(
    "The document text...",
    prompt="title: Using Prompts in RAG | text: "
)
    • Sans titre :
    doc_embedding = model.encode(
    "The document text...",
    prompt="title: none | text: "
)

Documentation complémentaire

Classification

La classification consiste à attribuer un texte à une ou plusieurs catégories ou étiquettes prédéfinies. Il s'agit de l'une des tâches les plus fondamentales du traitement du langage naturel (TLN).

Le routage des demandes d'assistance client est une application pratique de la classification de texte. Ce processus redirige automatiquement les demandes des clients vers le service approprié, ce qui permet de gagner du temps et de réduire le travail manuel.

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

Apprentissage de représentation Matriochka (MRL)

EmbeddingGemma utilise MRL pour fournir plusieurs tailles d'intégration à partir d'un seul modèle. Il s'agit d'une méthode d'entraînement intelligente qui crée un embedding unique de haute qualité, où les informations les plus importantes sont concentrées au début du vecteur.

Cela signifie que vous pouvez obtenir un embedding plus petit, mais toujours très utile, en prenant simplement les N premières dimensions de l'embedding complet. Le stockage et le traitement d'embeddings plus petits et tronqués sont beaucoup moins coûteux, mais cette efficacité se fait au détriment de la qualité potentiellement inférieure des embeddings. MRL vous permet de choisir l'équilibre optimal entre cette vitesse et cette précision en fonction des besoins spécifiques de votre application.

Utilisons trois mots ["apple", "banana", "car"] et créons des embeddings simplifiés pour voir comment fonctionne 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

Désormais, vous n'avez pas besoin d'un nouveau modèle pour une application plus rapide. Il vous suffit de tronquer les embeddings complets aux 512 premières dimensions. Pour des résultats optimaux, il est également recommandé de définir normalize_embeddings=True, qui met les vecteurs à une longueur unitaire 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

Dans les environnements extrêmement contraints, vous pouvez raccourcir davantage les embeddings à 256 dimensions. Vous pouvez également utiliser le produit scalaire plus efficace pour les calculs de similarité au lieu de la similarité cosinus standard.

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

Résumé et prochaines étapes

Vous êtes maintenant en mesure de générer des embeddings textuels de haute qualité à l'aide d'EmbeddingGemma et de la bibliothèque Sentence Transformers. Appliquez ces compétences pour créer des fonctionnalités puissantes telles que la similarité sémantique, la classification de texte et les systèmes de génération augmentée par récupération (RAG), et continuez à explorer les possibilités offertes par les modèles Gemma.

Consultez ensuite les documents suivants :