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 w GitHubie

EmbeddingGemma to lekki, otwarty model do tworzenia osadzeń, który został zaprojektowany z myślą o szybkim i wysokiej jakości wyszukiwaniu 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 połączenia z internetem.

Konfiguracja

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

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

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

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 z 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 model EmbeddingGemma przetworzy 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 nanieśli 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 są do siebie różne zdania pod względem semantycznym. Poniżej znajdziesz przykłady z wysokim, średnim i niskim wynikiem podobieństwa.

  • Wysokie podobieństwo:

    • Zdanie A: „The chef prepared a delicious meal for the guests.”
    • 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: „Bardzo interesuje się sztuczną inteligencją”.
    • Uzasadnienie: zdania są powiązane, ponieważ uczenie maszynowe jest poddziedziną sztucznej inteligencji. Mówią jednak o różnych osobach i mają różny poziom zaangażowania (ekspert vs. osoba zainteresowana).

  • 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 EmbeddingGemma

Aby wygenerować najlepsze osadzanie za pomocą 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. Daje to 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 prompta domyślnego
    Jeśli nie określisz ani prompt, ani prompt_name, system automatycznie użyje prompta ustawionego jako default_prompt_name. Jeśli nie ustawiono prompta 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 rozszerzone przez wyszukiwanie w zapisanych informacjach (RAG)

W przypadku systemów RAG użyj tych wartości prompt_name, aby utworzyć specjalistyczne wektory osadzeń 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 działu obsługi klienta. Ten proces automatycznie kieruje zapytania klientów do odpowiedniego działu, co pozwala zaoszczędzić czas i zmniejszyć nakład pracy ręcznej.

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)

Model 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 wymiary 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ż mieć nowego modelu. Po prostu skróć pełne osadzenia do pierwszych 512 wymiarów. Aby uzyskać optymalne wyniki, zalecamy też ustawienie normalize_embeddings=True, które 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 osadzanie tekstu za pomocą EmbeddingGemma i biblioteki Sentence Transformers. Wykorzystaj te umiejętności do tworzenia zaawansowanych funkcji, takich jak podobieństwo semantyczne, klasyfikacja tekstu i systemy Retrieval-Augmented Generation (RAG), oraz odkrywaj możliwości modeli Gemma.

Zapoznaj się z tymi dokumentami: