Как создавать векторные представления с помощью Sentence Transformers

Посмотреть на ai.google.dev Запустить в Google Colab Запустить в Kaggle Открыть в Vertex AI Посмотреть исходный код на GitHub

EmbeddingGemma – это легкая открытая модель для создания встраиваемых объектов, предназначенная для быстрого и качественного поиска на обычных устройствах, таких как мобильные телефоны. Модель содержит всего 308 миллионов параметров, поэтому она достаточно эффективна для использования передовых методов ИИ, таких как генерация с расширенным поиском (RAG), непосредственно на вашем локальном устройстве без подключения к интернету.

Настройка

Прежде чем начать работу с этим руководством, выполните следующие действия:

  • Чтобы получить доступ к Gemma, войдите в аккаунт Hugging Face и выберите Acknowledge license (Принять лицензию) для модели Gemma.
  • Сгенерируйте токен доступа Hugging Face и используйте его для входа в Colab.

Этот блокнот будет выполняться на центральном или графическом процессоре.

Как установить пакеты Python

Установите библиотеки, необходимые для запуска модели EmbeddingGemma и создания встраиваний. Sentence Transformers – это фреймворк на языке Python для встраивания текста и изображений. Дополнительную информацию можно найти в документации по Sentence Transformers.

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

После принятия лицензии вам понадобится действительный токен Hugging Face, чтобы получить доступ к модели.

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

Загрузить модель

Используйте библиотеки sentence-transformers, чтобы создать экземпляр класса модели с помощью 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

Создание встраивания

Векторное представление – это числовое представление текста, например слова или предложения, которое отражает его семантическое значение. По сути, это список чисел (вектор), который позволяет компьютерам понимать отношения между словами и их контекст.

Давайте посмотрим, как EmbeddingGemma обработает три разных слова ["apple", "banana", "car"].

EmbeddingGemma обучена на большом количестве текстов и знает, как связаны слова и понятия.

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

Модель выводит числовой вектор для каждого предложения. Векторы на самом деле очень длинные (768), но для простоты они представлены с несколькими измерениями.

Важны не сами числа, а расстояние между векторами. Если бы мы нанесли эти векторы на многомерное пространство, то векторы для apple и banana были бы очень близки друг к другу. А вектор для car будет далеко от двух других.

Как определить сходство

В этом разделе мы используем встраивание, чтобы определить, насколько семантически похожи разные предложения. Ниже приведены примеры с высоким, средним и низким показателями сходства.

  • Высокое сходство:

    • Предложение А: "Шеф-повар приготовил для гостей вкусное блюдо".
    • Предложение Б: "Шеф-повар приготовил для посетителей вкусный ужин".
    • Обоснование. В обоих предложениях описывается одно и то же событие, но разными словами и грамматическими конструкциями (активный и пассивный залог). Они передают один и тот же основной смысл.

  • Среднее сходство:

    • Предложение А: "Она эксперт в области машинного обучения".
    • Предложение Б: "Он очень интересуется искусственным интеллектом".
    • Обоснование: предложения связаны, поскольку машинное обучение – это область искусственного интеллекта. Однако в них говорится о разных людях с разным уровнем вовлеченности (эксперт и заинтересованный пользователь).

  • Низкое сходство:

    • Предложение А: "Сегодня в Токио солнечная погода".
    • Предложение Б: "Мне нужно купить продукты на неделю".
    • Обоснование. В этих двух предложениях речь идет о совершенно разных вещах, и они не связаны между собой по смыслу.
# 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

Как использовать запросы с EmbeddingGemma

Чтобы получить наилучшие встраивания с помощью EmbeddingGemma, добавьте в начало входного текста "инструкцию" или "задачу". Эти запросы оптимизируют встраивания для определенных задач, таких как поиск документов или ответы на вопросы, и помогают модели различать разные типы входных данных, например поисковый запрос и документ.

Как применять запросы

Применить запрос во время вывода можно тремя способами.

  1. Используйте аргумент prompt
    Передайте полную строку запроса непосредственно методу encode. Это позволяет точно управлять настройками.

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

  2. Используя аргумент prompt_name
    Выберите готовый запрос по его названию. Эти подсказки загружаются из конфигурации модели или во время ее инициализации.

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

  3. Использование запроса по умолчанию
    Если вы не укажете ни prompt, ни prompt_name, система автоматически использует запрос, заданный как default_prompt_name. Если запрос по умолчанию не задан, то запрос не применяется.

    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

Пример использования: генерация с дополненной выборкой (RAG)

Для систем RAG используйте следующие значения prompt_name, чтобы создавать специализированные встраивания для запросов и документов:

  • Для запросов: используйте prompt_name="Retrieval-query".

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

  • Для документов: используйте prompt_name="Retrieval-document". Чтобы ещё больше улучшить встраивание документов, вы также можете добавить заголовок, используя аргумент prompt напрямую:

    • С заголовком:
    doc_embedding = model.encode(
    "The document text...",
    prompt="title: Using Prompts in RAG | text: "
)
    • Без названия:
    doc_embedding = model.encode(
    "The document text...",
    prompt="title: none | text: "
)

Дополнительные материалы

Классификация

Классификация – это задача по отнесению фрагмента текста к одной или нескольким заранее определенным категориям или ярлыкам. Это одна из основных задач в области обработки естественного языка (NLP).

Практическое применение классификации текста – маршрутизация запросов в службу поддержки. Этот процесс автоматически направляет запросы клиентов в нужный отдел, экономя время и сокращая объем ручной работы.

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)

Встраивание Gemma использует MRL, чтобы предоставлять несколько размеров встраивания из одной модели. Это эффективный метод обучения, который позволяет создать одно высококачественное встраивание, в котором наиболее важная информация сосредоточена в начале вектора.

Это означает, что вы можете получить меньший, но все ещё очень полезный вектор, просто взяв первые N измерений полного вектора. Хранить и обрабатывать усеченные встраивания значительно дешевле и быстрее, но это может привести к снижению качества. MRL позволяет выбрать оптимальный баланс между скоростью и точностью в зависимости от потребностей вашего приложения.

Давайте рассмотрим, как работает MRL, на примере трех слов ["apple", "banana", "car"] и упрощенных встраиваний.

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

Теперь для более быстрого применения не нужна новая модель. Просто обрежьте полные встраивания до первых 512 измерений. Для достижения оптимальных результатов также рекомендуется задать значение normalize_embeddings=True, которое масштабирует векторы до единичной длины 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

В условиях очень ограниченных ресурсов можно сократить количество измерений до 256. Для расчета сходства можно использовать более эффективный скалярное произведение вместо стандартного косинусного сходства.

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

Выводы и дальнейшие действия

Теперь вы можете создавать высококачественные текстовые встраивания с помощью EmbeddingGemma и библиотеки Sentence Transformers. Используйте эти навыки, чтобы создавать мощные функции, такие как семантическое сходство, классификация текста и системы генерации с дополненной выборкой (RAG), и продолжайте изучать возможности моделей Gemma.

Рекомендуем ознакомиться со следующими документами: