Generowanie wektorów dystrybucyjnych za pomocą Sentence Transformers

Wyświetl na ai.google.dev Uruchom w Google Colab Uruchom w Kaggle Otwórz w Vertex AI Wyświetl źródło na GitHubie

EmbeddingGemma to lekki, otwarty model wektorów dystrybucyjnych zaprojektowany z myślą o szybkim wyszukiwaniu wysokiej jakości na urządzeniach codziennego użytku, takich jak telefony komórkowe. Ma tylko 308 milionów parametrów, więc jest wystarczająco wydajny, aby uruchamiać zaawansowane techniki AI, takie jak Retrieval Augmented Generation (RAG), bezpośrednio na urządzeniu lokalnym bez konieczności połączenia z internetem.

Konfiguracja

Zanim rozpoczniesz wykonywanie zadań z tego samouczka, wykonaj te czynności:

  • Aby uzyskać dostęp do Gemy, zaloguj się w Hugging Face i wybierz Potwierdź licencję dla modelu Gemy.
  • Wygeneruj token dostępu Hugging Face i użyj go do zalogowania się w Colab.

Ten notatnik będzie działać na procesorze lub procesorze graficznym.

Instalowanie pakietów Pythona

Zainstaluj biblioteki wymagane do uruchomienia modelu EmbeddingGemma i generowania wektorów dystrybucyjnych. Sentence Transformers to platforma Pythona do tworzenia reprezentacji właściwościowych tekstu i obrazów. Więcej informacji znajdziesz w dokumentacji Sentence Transformers.

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

Po zaakceptowaniu licencji musisz mieć ważny token Hugging Face, aby uzyskać dostęp do modelu.

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

Wczytaj model

Użyj bibliotek sentence-transformers, aby utworzyć instancję klasy modelu za pomocą 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

Generowanie wektora dystrybucyjnego

Wektor dystrybucyjny to numeryczna reprezentacja tekstu, np. słowa lub zdania, która odzwierciedla jego znaczenie semantyczne. Jest to lista liczb (wektor), która umożliwia komputerom zrozumienie relacji i kontekstu słów.

Zobaczmy, jak EmbeddingGemma przetwarza 3 różne słowa["apple", "banana", "car"].

Model EmbeddingGemma został wytrenowany na ogromnych ilościach tekstu i nauczył się relacji między słowami i pojęciami.

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

Model zwraca wektor liczbowy dla każdego zdania. Rzeczywiste wektory są bardzo długie (768), ale dla uproszczenia przedstawiamy je w kilku wymiarach.

Kluczem nie są poszczególne liczby, ale odległość między wektorami. Gdybyśmy wykreślili te wektory w przestrzeni wielowymiarowej, wektory dla apple i banana byłyby bardzo blisko siebie. Wektor dla car będzie znacznie oddalony od pozostałych dwóch.

Określanie podobieństwa

W tej sekcji używamy osadzania, aby określić, jak podobne semantycznie są różne zdania. Poniżej przedstawiamy przykłady z wysokimi, średnimi i niskimi wynikami podobieństwa.

  • Wysokie podobieństwo:

    • Zdanie A: „Szef kuchni przygotował dla gości pyszny posiłek”.
    • Zdanie B: „Pyszny obiad został ugotowany przez szefa kuchni dla gości”.
    • Uzasadnienie: oba zdania opisują to samo zdarzenie, ale używają różnych słów i struktur gramatycznych (strona czynna i bierna). Przekazują to samo podstawowe znaczenie.

  • Średnie podobieństwo:

    • Zdanie A: „She is an expert in machine learning”.
    • Zdanie B: „He has a deep interest in artificial intelligence”.
    • Uzasadnienie: zdania są powiązane, ponieważ uczenie maszynowe jest poddziedziną sztucznej inteligencji. Mówią jednak o różnych osobach o różnym poziomie zaangażowania (eksperci i osoby zainteresowane).

  • Niskie podobieństwo:

    • Zdanie A: „Dziś w Tokio jest słonecznie”.
    • Zdanie B: „Muszę kupić produkty spożywcze na ten tydzień”.
    • Uzasadnienie: te 2 zdania dotyczą zupełnie różnych tematów i nie mają ze sobą żadnego związku semantycznego.
# 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

Używanie promptów z modelem EmbeddingGemma

Aby wygenerować najlepsze osadzanie za pomocą modelu EmbeddingGemma, na początku tekstu wejściowego dodaj „prompt z instrukcjami” lub „zadanie”. Te prompty optymalizują osadzanie pod kątem konkretnych zadań, takich jak wyszukiwanie dokumentów czy odpowiadanie na pytania, i pomagają modelowi rozróżniać różne typy danych wejściowych, np. zapytanie wyszukiwania od dokumentu.

Jak stosować prompty

Podczas wnioskowania możesz zastosować prompt na 3 sposoby.

  1. Używanie argumentu prompt
    Przekaż pełny ciąg prompta bezpośrednio do metody encode. Dzięki temu masz precyzyjną kontrolę.

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

  2. Używanie argumentu prompt_name
    Wybierz gotowy prompt według nazwy. Te prompty są wczytywane z konfiguracji modelu lub podczas jego inicjowania.

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

  3. Używanie promptu domyślnego
    Jeśli nie określisz ani prompt, ani prompt_name, system automatycznie użyje promptu ustawionego jako default_prompt_name. Jeśli nie ustawiono promptu domyślnego, nie zostanie zastosowany żaden prompt.

    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

Przypadek użycia: generowanie wspomagane wyszukiwaniem (RAG)

W przypadku systemów RAG użyj tych wartości prompt_name, aby utworzyć specjalistyczne osadzanie dla zapytań i dokumentów:

  • W przypadku zapytań: użyj prompt_name="Retrieval-query".

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

  • W przypadku dokumentów: użyj prompt_name="Retrieval-document". Aby jeszcze bardziej ulepszyć osadzanie dokumentów, możesz też dodać tytuł, używając bezpośrednio argumentu prompt:

    • Z tytułem:
    doc_embedding = model.encode(
    "The document text...",
    prompt="title: Using Prompts in RAG | text: "
)
    • Bez tytułu:
    doc_embedding = model.encode(
    "The document text...",
    prompt="title: none | text: "
)

Więcej informacji

Klasyfikacja

Klasyfikacja to przypisywanie fragmentu tekstu do co najmniej 1 wstępnie zdefiniowanej kategorii lub etykiety. Jest to jedno z najważniejszych zadań w przetwarzaniu języka naturalnego (NLP).

Praktycznym zastosowaniem klasyfikacji tekstu jest kierowanie zgłoszeń do zespołu obsługi klienta. Ten proces automatycznie kieruje zapytania klientów do odpowiedniego działu, co pozwala zaoszczędzić czas i ograniczyć zakres czynności wykonywanych ręcznie.

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

Matryoshka Representation Learning (MRL)

EmbeddingGemma wykorzystuje MRL, aby udostępniać wiele rozmiarów wektorów dystrybucyjnych z jednego modelu. Jest to sprytna metoda trenowania, która tworzy jeden wektor dystrybucyjny wysokiej jakości, w którym najważniejsze informacje są skoncentrowane na początku wektora.

Oznacza to, że możesz uzyskać mniejszy, ale nadal bardzo przydatny wektor, po prostu biorąc pierwsze N wymiarów pełnego wektora. Przechowywanie mniejszych, skróconych wektorów osadzania jest znacznie tańsze, a ich przetwarzanie szybsze, ale ta wydajność wiąże się z potencjalnie niższą jakością wektorów osadzania. MRL umożliwia wybór optymalnej równowagi między szybkością a dokładnością w zależności od konkretnych potrzeb aplikacji.

Użyjmy 3 słów ["apple", "banana", "car"] i utwórzmy uproszczone osadzanie, aby zobaczyć, jak działa 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

Aby przyspieszyć proces składania wniosku, nie musisz już tworzyć nowego modelu. Po prostu skróć pełne wektory do pierwszych 512 wymiarów. Aby uzyskać optymalne wyniki, zalecamy też ustawienie wartości normalize_embeddings=True, która skaluje wektory do długości jednostkowej 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

W środowiskach o bardzo ograniczonych zasobach możesz dodatkowo skrócić wektory do 256 wymiarów. Do obliczania podobieństwa możesz też używać bardziej wydajnego iloczynu skalarnego zamiast standardowego podobieństwa cosinusowego.

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

Podsumowanie i dalsze kroki

Możesz teraz generować wysokiej jakości wektory dystrybucyjne tekstu za pomocą modelu EmbeddingGemma i biblioteki Sentence Transformers. Wykorzystaj te umiejętności do tworzenia zaawansowanych funkcji, takich jak podobieństwo semantyczne, klasyfikacja tekstu i systemy generowania wspomaganego wyszukiwaniem (RAG), oraz odkrywaj możliwości modeli Gemma.

Zapoznaj się z tymi dokumentami: