Встраивание текста и поиск с помощью сходства предложений в масштабе с Amazon SageMaker JumpStart

Максимальное использование Amazon SageMaker JumpStart для встраивания текста и поиска на основе сходства предложений

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

Прежде чем создавать интеллектуальные приложения из вложений, предприятиям и организациям приходилось встраивать свои существующие документы, что может быть дорого и технически сложно. Amazon SageMaker JumpStart – это центр машинного обучения (ML), который помогает ускорить этот процесс. С помощью SageMaker JumpStart вы можете получить доступ к предобученным передовым моделям вложения текста от разных поставщиков моделей, включая Hugging Face, AI 21 Labs, Cohere и Meta AI. Вы можете безупречно развернуть эти модели в производстве с помощью пользовательского интерфейса или SDK JumpStart SageMaker. Кроме того, ни одни из ваших данных не используются для обучения основных моделей. Поскольку все данные зашифрованы и не покидают собственную виртуальную частную сеть, вы можете быть уверены, что ваши данные остаются конфиденциальными и неприступными.

В этой статье мы покажем, как использовать SageMaker Python SDK для вложения текста и сравнения предложений. Сравнение предложений включает оценку сходства между двумя кусками текста после их преобразования в вложения LLM, что является основным шагом для таких приложений, как Retrieval Augmented Generation (RAG). Мы показываем, как сделать следующее:

  • Выполнять вывод на модель вложения текста, развернутую с помощью SageMaker JumpStart
  • Находить ближайших соседей для входного предложения с использованием собственного набора данных
  • Выполнять пакетное преобразование на больших документах для минимизации затрат

Весь код доступен на GitHub.

Развернуть модель вложения текста через SageMaker JumpStart

Для размещения модели на Amazon SageMaker, первым шагом является настройка и аутентификация использования сервисов AWS. В Amazon SageMaker Studio мы используем роль выполнения, связанную с экземпляром блокнота. См. следующий код:

import sagemaker, boto3, jsonfrom sagemaker.session import Sessionsagemaker_session = Session()aws_role = sagemaker_session.get_caller_identity_arn()aws_region = boto3.Session().region_namesess = sagemaker.Session()

На Hugging Face Massive Text Embedding Benchmark (MTEB) предоставляется в качестве таблицы лидеров для различных задач вложения текста. В настоящее время он предоставляет 129 наборов бенчмарков на 8 различных задач на 113 языках. Лучшие модели вложения текста из таблицы лидеров MTEB стали доступны в SageMaker JumpStart, включая bge, gte, e5 и многие другие. В этой статье мы используем huggingface-sentencesimilarity-bge-large-en в качестве примера. Мы можем использовать SDK SageMaker для развертывания этой передовой модели вложения текста:

from sagemaker.jumpstart.model import JumpStartModelmodel_id = "huggingface-sentencesimilarity-bge-large-en"text_embedding_model = JumpStartModel(model_id=model_id)predictor = text_embedding_model.deploy()

Запрос модели вложения текста

Давайте подробнее рассмотрим запрос модели вложения текста.

Текст в вложение

Если вы уже развернули конечную точку SageMaker ранее, предиктор можно восстановить следующим образом:

from sagemaker.predictor import Predictorfrom sagemaker.deserializers import JSONDeserializerfrom sagemaker.serializers import IdentitySerializerpredictor = Predictor(    endpoint_name=<YOUR_ENDPOINT_NAME>,    deserializer=JSONDeserializer(),    serializer=IdentitySerializer(),)predictor.content_type = "application/x-text"

После успешного развертывания модели вы можете запросить конечную точку с пакетом входных текстов внутри JSON-нагрузки:

sentences = [    # Домашние животные    "Твоя собака такая милая.",    "Какая милая твоя собака!",    "У тебя такая милая собака!",    # Города    "Сидней - это место, где я работаю.",    "Я работаю в Сиднее.",    # Цвет    "Какой цвет тебе нравится больше всего?",    "Какой твой самый любимый цвет?",]predictor.predict(json.dumps(sentences).encode('utf-8'))

Корреляция векторных представлений этих предложений отображена на следующей диаграмме.

correlation_heat_map

Как показано на предыдущей диаграмме, одинаковые темы сильно коррелируют между собой, включая Домашние животные, Города и Цвет; разные темы значительно отличаются. Это говорит о том, что векторное представление, сгенерированное моделью LLMs (в данном случае, bge), может точно представлять семантическую информацию.

В этом сообщении мы использовали предыдущий пример и сравнили задержку для разных моделей представления предложений, доступных в SageMaker JumpStart. Задержка представляет собой время от момента отправки пользователем запроса до момента завершения приложением выполнения запроса. Числа в следующей таблице представляют среднюю задержку для общего числа 100 запросов с использованием того же пакета входных текстов на экземплярах ml.g5.2xlarge и ml.c6i.xlarge.

Модель Средняя задержка для g5.2xlarge (мс) Средняя задержка для c6i.xlarge (мс) Поддержка языка
all-MiniLM-L6-v2 19.5 27.9 Английский
BGE Base En 21.2 114 Английский
BGE Small En 28.3 45.6 Английский
BGE Large En 34.7 337 Английский
Multilingual E5 Base 22.1 118 Мультиязычный
Multilingual E5 Large 39.8 360 Мультиязычный
E5 Base 25.6 117 Английский
E5 Base V2 25.2 123 Английский
E5 Large 32.2 339 Английский
E5 Large V2 32.5 331 Английский
GTE Base 22.2 112 Английский
GTE Small 19.7 46 Английский
GTE Large 39.7 347 Английский

Получение ближайших соседей

Развернутая модель от SageMaker JumpStart также может облегчить процесс определения ближайших соседей для запросов в корпусе. Когда предоставлены запросы и корпус, модель будет выдавать corpus_id, который указывает на позицию соответствующей записи корпуса в списке входных данных, и оценку, указывающую на степень близости к запросу. Она использует следующие параметры:

  • corpus – Предоставляет список входных данных, на основе которого следует найти ближайшего соседа
  • queries – Предоставляет список входных данных, для которых нужно найти ближайшего соседа в корпусе
  • top_k – Количество ближайших соседей, которых следует найти в корпусе
  • mode – Установлено как nn_corpus для получения ближайших соседей для входных запросов внутри корпуса

Рассмотрите следующий код:

corpus = [    "Amazon SageMaker – это полностью управляемый сервис для подготовки данных и создания, обучения и развертывания моделей машинного обучения (ML) для любых случаев использования с полностью управляемой инфраструктурой, инструментами и рабочими процессами.",    "Amazon SageMaker хранит код на томах хранения ML, защищенных группами безопасности и при необходимости зашифрованных в покое.",    "Amazon SageMaker предоставляет полный рабочий процесс от начала до конца, но вы также можете продолжать использовать свои существующие инструменты с SageMaker. Вы можете легко передавать результаты каждого этапа внутри и вне SageMaker в соответствии с бизнес-требованиями."]queries = [    "Что такое Amazon SageMaker?",    "Как Amazon SageMaker защищает мой код?",    "Что делать, если у меня есть собственная среда блокнота, обучения или хостинга в собственной рабочей среде?"]payload_nearest_neighbor = {"corpus": corpus, "queries": queries, "top_k": 3, "mode": "nn_corpus"}query_response = predictor.predict(payload_nearest_neighbor)

Мы получаем следующий вывод:

[    [        {'corpus_id': 0, 'score': 0.8992230892181396},        {'corpus_id': 2, 'score': 0.8664969205856323},        {'corpus_id': 1, 'score': 0.8456423282623291}    ],    [        {'corpus_id': 1, 'score': 0.8919335603713989},        {'corpus_id': 0, 'score': 0.840064525604248},        {'corpus_id': 2, 'score': 0.8145401477813721}    ],    [        {'corpus_id': 2, 'score': 0.7712811231613159},        {'corpus_id': 1, 'score': 0.7564010620117188},        {'corpus_id': 0, 'score': 0.7525666356086731}    ]]

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

Мы также взяли предыдущий пример и сравнили задержку для различных моделей встраивания предложений, которые в данный момент доступны в SageMaker JumpStart. Числа в следующей таблице представляют среднюю задержку при общем количестве 100 запросов с использованием того же набора данных на экземплярах ml.g5.2xlarge и ml.c6i.xlarge.

Модель Средняя задержка g5.2xlarge (мс) Средняя задержка c6i.xlarge (мс) Поддержка языка
all-MiniLM-L6-v2 21.7 69.1 Английский
BGE Base En 29.1 372 Английский
BGE Small En 29.2 124 Английский
BGE Large En 47.2 1240 Английский
Multilingual E5 Base 30 389 Мультиязычный
Multilingual E5 Large 47.1 1380 Мультиязычный
E5 Base 30.4 373 Английский
E5 Base V2 31 409 Английский
E5 Large 45.9 1230 Английский
E5 Large V2 49.6 1220 Английский
GTE Base 30.3 375 Английский
GTE Small 28.5 129 Английский
GTE Large 46.6 1320 Английский

Получите ближайших соседей на большом наборе данных

При отправке запросов к точке вызова SageMaker ограничение на размер полезной нагрузки составляет примерно 5 МБ, а время ожидания запроса установлено на 1 минуту. Если размер корпуса превышает эти ограничения, можно использовать задание обучения SageMaker, которое генерирует вложения для вашего большого набора данных и сохраняет их вместе с моделью внутри точки вызова SageMaker. Поэтому их не нужно передавать в качестве части полезной нагрузки вызова. Процесс поиска ближайших соседей осуществляется с помощью SentenceTransformer и его служебной функции. Ближайший сосед основывается на косинусном сходстве между встраиванием вводной предложения и предварительно вычисленными встраиваниями предложений во время задания обучения.

В следующем примере мы получаем и подготавливаем набор данных Amazon_SageMaker_FAQs, чтобы использовать его для поиска ближайшего соседа для входного вопроса:

!aws s3 cp s3://jumpstart-cache-prod-us-west-2/training-datasets/Amazon_SageMaker_FAQs/Amazon_SageMaker_FAQs.csv Amazon_SageMaker_FAQs.csvimport pandas as pddata = pd.read_csv("Amazon_SageMaker_FAQs.csv", names=["Questions", "Answers"])data["id"] = data.indexdata_req = data[["id", "Answers"]]data_req.to_csv("data.csv", index=False, header=False)output_bucket = sess.default_bucket()output_prefix = "jumpstart-example-ss-training"s3_output_location = f"s3://{output_bucket}/{output_prefix}/output"training_dataset_s3_path = f"s3://{output_bucket}/{output_prefix}/data/data.csv"!aws s3 cp data.csv {training_dataset_s3_path}

Для алгоритмо-специфичных гиперпараметров обучения можно извлечь или перезаписать SDK SageMaker:

from sagemaker import hyperparametershyperparameters = hyperparameters.retrieve_default(model_id=model_id, model_version = "*")hyperparameters["batch_size"] = "64"print(hyperparameters)>>> {'max_seq_length': 'None', 'batch_size': '64', 'store_text_with_embedding': 'True'}

Обучение SageMaker состоит из двух шагов: создание объекта оценщика и запуск задания обучения. Результатом является модель, упакованная с встраиваниями вашего большого набора данных, используемого в качестве обучающих данных, которую можно развернуть для вывода ближайшего соседа для любого входного предложения. См. следующий код:

from sagemaker.jumpstart.estimator import JumpStartEstimatorestimator = JumpStartEstimator(    model_id=model_id,    hyperparameters=hyperparameters,    output_path=s3_output_location)estimator.fit(    {"training": f"s3://{output_bucket}/{output_prefix}/data"})predictor = estimator.deploy()

Синтаксис запроса для преобразования текста в встраивания остается прежним. Однако код для получения ближайшего соседа можно упростить следующим образом:

payload_nearest_neighbour = {    "queries": ["Поддерживается ли R в Amazon SageMaker?"],    "top_k": 1,    "mode": "nn_train_data",}response = predictor.predict(payload_nearest_neighbour)>>> [[{'id': '9', 'score': 0.9240573048591614}]]data["Answers"].iloc[int(response[0][0]["id"])]>>> "Да, R поддерживается в Amazon SageMaker. Вы можете использовать R внутри экземпляров блокнотов SageMaker, включающих предустановленное ядро R и библиотеку reticulate. Reticulate предлагает интерфейс R для Python SDK Amazon SageMaker, позволяющий специалистам по машинному обучению создавать, обучать, настраивать и развертывать модели R."

Мы также можем обратиться к конечной точке с вопросами в наборе данных Amazon_SageMaker_FAQs и сравнить, сколько правильных соответствующих ответов возвращается. В следующем примере мы измеряем точность топ-3, учитывая, что может быть похожие пары вопросов и ответов. Это означает, что если правильный ответ возвращается как одно из трех лучших ответов, то это считается правильным запросом.

total_correct_answers = 0for i in range(len(data)):    question = data["Questions"].iloc[i]    payload_nearest_neighbor = {        "queries": [question],        "top_k": 3,        "mode": "nn_train_data",    }    response = predictor.predict(payload_nearest_neighbor)    response_ids = [int(res["id"]) for res in response[0]]    if i in response_ids:        total_correct_answers += 1    else:        pred_answer = [data["Answers"].iloc[response_id] for response_id in response_ids]print(total_correct_answers*100/len(data))>>> 81.16883116883118

Запустите пакетное преобразование для получения эмбеддингов на больших наборах данных

Для предприятий и организаций с большим объемом исторических документов, превышающих память одного экземпляра конечной точки, можно использовать пакетное преобразование SageMaker для экономии затрат. При запуске задания пакетного преобразования SageMaker запускает необходимые вычислительные ресурсы для обработки данных. Во время задания SageMaker автоматически предоставляет и управляет вычислительными ресурсами. По завершении задания пакетного преобразования эти ресурсы автоматически очищаются, что минимизирует затраты. Разделяя большой набор данных на более маленькие части и используя больше экземпляров, вы можете масштабировать вычисления для более быстрого вывода схожей стоимости, не управляя инфраструктурой. Максимальная полезная нагрузка для пакетного преобразования составляет 100 МБ, а время ожидания – 1 час.

Формат ввода для нашей задачи пакетного преобразования – это файл JSONL с записями в виде строки JSON, содержащий id и text_inputs. Смотрите следующий код:

test_data_file_name = "test.jsonl"test_data = []for i in range(len(data)):    answer = data.loc[i, "Answers"]    payload = {"id": i, "text_inputs": answer}    test_data.append(payload)with open(test_data_file_name, "w") as outfile:    for entry in test_data:        outfile.write(f"{json.dumps(entry)}\n")s3 = boto3.client("s3")s3.upload_file(test_data_file_name, output_bucket, f"{output_prefix}/batch_input/test.jsonl")

Когда данные готовы в Amazon Simple Storage Service (Amazon S3), вы можете создать объект пакетного преобразования из модели SageMaker JumpStart, которая запускает задание преобразования:

s3_input_data_path = f"s3://{output_bucket}/{output_prefix}/batch_input/"s3_output_data_path = f"s3://{output_bucket}/{output_prefix}/batch_output/"batch_transformer = text_embedding_model.transformer(    instance_count=1,    instance_type="ml.p3.2xlarge",    output_path=s3_output_data_path,    assemble_with="Line",    accept="text/csv",    max_payload=1,)batch_transformer.transform(    s3_input_data_path,    content_type="application/jsonlines",    split_type="Line")batch_transformer.wait()

После завершения задания пакетного преобразования результат можно скачать из Amazon S3:

s3 = boto3.client("s3")s3.download_file(    output_bucket, output_prefix + "/batch_output/" + "test.jsonl.out", "predict.jsonl")with open("predict.jsonl", "r") as json_file:    json_list = list(json_file)

Заключение

SageMaker JumpStart предоставляет простой способ использования современных моделей основных языковых моделей для встраивания текста и семантического поиска. С помощью пользовательского интерфейса или всего нескольких строк кода вы можете развернуть высокоточную модель встраивания текста и находить семантические соответствия в больших наборах данных масштабно и экономично. SageMaker JumpStart устраняет преграды к внедрению семантического поиска, предоставляя мгновенный доступ к передовым моделям, таким как модели, прошедшие проверку на доске MTEB. Бизнес и разработчики могут быстрее создавать интеллектуальные системы поиска и рекомендаций.

В этой статье показано, как находить семантически схожие вопросы и ответы, что может быть применено к случаям использования RAG, рекомендациям и персонализации, многоязычным переводам и другим. Благодаря постоянному развитию языковых моделей и простоте SageMaker JumpStart, большее количество организаций может внедрять возможности генеративного ИИ в свои продукты. В качестве следующего шага вы можете использовать модели встраивания текста из SageMaker JumpStart на своем собственном наборе данных для тестирования и оценки результатов для вашего случая использования RAG.