Incrustación de texto y recuperación de similitudes de oraciones a escala con Amazon SageMaker JumpStart

Los vectores de texto o incrustaciones son representaciones vectoriales numéricas de texto generadas por modelos de lenguaje grandes (LLM). Después de que los LLM estén completamente entrenados previamente en un gran conjunto de datos o ajustados a partir de diferentes tareas, incluida la finalización de texto, la respuesta a preguntas y las traducciones, las incrustaciones de texto capturan información semántica del texto de entrada. La incrustación de texto hace posible diferentes aplicaciones posteriores, incluida la búsqueda de similitudes, la recuperación de información, las recomendaciones y la personalización, las traducciones multilingües y más.

Antes de que se pudieran crear aplicaciones inteligentes a partir de incorporaciones, las empresas y organizaciones tenían que incorporar sus documentos existentes, lo que puede resultar costoso y técnicamente complicado. Inicio rápido de Amazon SageMaker es un centro de aprendizaje automático (ML) que ayuda a acelerar este viaje. Con SageMaker JumpStart, puede acceder a modelos de incrustación de texto de vanguardia y previamente entrenados de varios proveedores de modelos, incluidos Hugging Face, AI 21 Labs, Cohere y Meta AI. Puede implementar estos modelos en producción sin problemas con la interfaz de usuario o SDK de SageMaker JumpStart. Además, ninguno de sus datos se utiliza para entrenar los modelos subyacentes. Debido a que todos los datos están cifrados y no salen de su propia VPC, puede confiar en que sus datos permanecerán privados y confidenciales.

En esta publicación, demostramos cómo usar el SDK de Python de SageMaker para incrustar texto y similitud de oraciones. La similitud de oraciones implica evaluar la semejanza entre dos fragmentos de texto después de que el LLM los convierte en incrustaciones, lo cual es un paso básico para aplicaciones como Retrieval Augmented Generation (RAG). Demostramos cómo hacer lo siguiente:

• Ejecutar inferencia en un modelo de incrustación de texto implementado desde SageMaker JumpStart
• Encuentre los vecinos más cercanos para una oración de entrada con su propio conjunto de datos
• Ejecute la transformación por lotes en documentos grandes para minimizar costos

Todo el código está disponible en GitHub.

Implementar un modelo de incrustación de texto a través de SageMaker JumpStart

Para alojar un modelo en Amazon SageMaker, el primer paso es configurar y autenticar el uso de los servicios de AWS. En Amazon SageMaker Estudio, utilizamos la función de ejecución asociada con la instancia del cuaderno. Vea el siguiente código:

import sagemaker, boto3, json
from sagemaker.session import Session
sagemaker_session = Session()
aws_role = sagemaker_session.get_caller_identity_arn()
aws_region = boto3.Session().region_name
sess = sagemaker.Session()

En Hugging Face, el punto de referencia de incrustación de texto masivo (MTEB) se proporciona como una tabla de clasificación para diversas tareas de incrustación de texto. Actualmente proporciona 129 conjuntos de datos de evaluación comparativa en 8 tareas diferentes en 113 idiomas. Los principales modelos de incrustación de texto de la tabla de clasificación MTEB están disponibles en SageMaker JumpStart, incluidos bge, gte, e5, y más. En esta publicación utilizamos huggingface-sentencesimilarity-bge-large-en como ejemplo. Podemos utilizar el SDK de SageMaker para implementar este modelo de incrustación de texto de última generación:

from sagemaker.jumpstart.model import JumpStartModel

model_id = "huggingface-sentencesimilarity-bge-large-en"
text_embedding_model = JumpStartModel(model_id=model_id)
predictor = text_embedding_model.deploy()

Consulta de modelo de incrustación de texto

Veamos la consulta del modelo de incrustación de texto con más detalle.

Texto para incrustar

Si ya ha implementado un punto final de SageMaker anteriormente, el predictor se puede restaurar de la siguiente manera:

from sagemaker.predictor import Predictor
from sagemaker.deserializers import JSONDeserializer
from sagemaker.serializers import IdentitySerializer

predictor = Predictor(
    endpoint_name=<YOUR_ENDPOINT_NAME>,
    deserializer=JSONDeserializer(),
    serializer=IdentitySerializer(),
)
predictor.content_type = "application/x-text"

Una vez que el modelo se haya implementado correctamente, puede consultar el punto final con un lote de textos de entrada dentro de una carga útil JSON:

sentences = [
    # Pets
    "Your dog is so cute.",
    "How cute your dog is!",
    "You have such a cute dog!",
    # Cities
    "Sydney is the place where I work.",
    "I work in Sydney.",
    # Color
    "What colour do you like the most?",
    "What is your favourite colour?",
]

predictor.predict(json.dumps(sentences).encode('utf-8'))

La correlación de las incrustaciones de estas oraciones se representa en la siguiente figura.

Como se muestra en la figura anterior, los mismos sujetos están altamente correlacionados dentro de sí mismos, incluidos Pets, Citiesy Color; Los diferentes temas son muy diferentes. Esto indica la incrustación generada por los LLM (en este caso, bge) puede representar la información semántica con precisión.

Para esta publicación, utilizamos el ejemplo anterior y comparamos la latencia entre diferentes modelos de incrustación de oraciones actualmente disponibles en SageMaker JumpStart. La latencia es la cantidad de tiempo desde el momento en que un usuario envía una solicitud hasta el momento en que la aplicación indica que la solicitud se ha completado. Los números de la siguiente tabla representan la latencia promedio para un total de 100 solicitudes utilizando el mismo lote de textos de entrada en el ml.g5.2xlarge y ml.c6i.xlarge instancias.

Consigue los vecinos más cercanos

El modelo implementado de SageMaker JumpStart también puede facilitar el proceso de identificación de los vecinos más cercanos a las consultas dentro del corpus. Cuando se le proporcionan consultas y un corpus, el modelo producirá el corpus_id, que denota la posición de la entrada del corpus relevante en la lista de corpus de entrada y una puntuación que indica el grado de proximidad a la consulta. Utiliza los siguientes parámetros:

• cuerpo – Proporciona la lista de entradas desde las cuales encontrar el vecino más cercano
• consultas – Proporciona la lista de entradas para encontrar el vecino más cercano del corpus.
• top_k – El número de vecinos más cercanos a encontrar en el corpus.
• modo – Establecer como nn_corpus para lograr que los vecinos más cercanos ingresen consultas dentro del corpus

Vea el siguiente código:

corpus = [
    "Amazon SageMaker is a fully managed service to prepare data and build, train, and deploy machine learning (ML) models for any use case with fully managed infrastructure, tools, and workflows.",
    "Amazon SageMaker stores code in ML storage volumes, secured by security groups and optionally encrypted at rest.",
    "Amazon SageMaker provides a full end-to-end workflow, but you can continue to use your existing tools with SageMaker. You can easily transfer the results of each stage in and out of SageMaker as your business requirements dictate."
]
queries = [
    "What is Amazon SageMaker?",
    "How does Amazon SageMaker secure my code?",
    "What if I have my own notebook, training, or hosting environment in my own business environment?"
]

payload_nearest_neighbor = {"corpus": corpus, "queries": queries, "top_k": 3, "mode": "nn_corpus"}
query_response = predictor.predict(payload_nearest_neighbor)

Obtenemos el siguiente resultado:

[
    [
        {'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}
    ]
]

Este resultado significa que la primera consulta es más similar al primer corpus, la segunda está más cerca del segundo corpus, y así sucesivamente. Esta es una coincidencia correcta en este ejemplo.

También tomamos el ejemplo anterior y comparamos la latencia entre diferentes modelos de incrustación de oraciones disponibles actualmente en SageMaker JumpStart. Los números en la siguiente tabla representan la latencia promedio para un total de 100 solicitudes usando la misma carga útil en el ml.g5.2xlarge y ml.c6i.xlarge instancias.

Obtenga los vecinos más cercanos en un gran conjunto de datos

Al realizar solicitudes al punto final de invocación de SageMaker, las cargas útiles se restringen a aproximadamente 5 MB y el tiempo de espera de la solicitud se establece en 1 minuto. Si el tamaño del corpus excede estos límites, puede utilizar un trabajo de entrenamiento de SageMaker, que genera incrustaciones para su gran conjunto de datos y las conserva junto con el modelo dentro del punto final de SageMaker. Por lo tanto, no es necesario pasarlos como parte de la carga útil de invocación. El proceso de búsqueda de vecinos más cercanos se realiza mediante Transformador de oraciones y su función de utilidad. El vecino más cercano se basa en la similitud del coseno entre la incrustación de oraciones de entrada y las incrustaciones de oraciones precalculadas durante el trabajo de entrenamiento.

En el siguiente ejemplo, buscamos y preparamos el Amazon_SageMaker_FAQs conjunto de datos para usarlo para encontrar el vecino más cercano a una pregunta de entrada:

!aws s3 cp s3://jumpstart-cache-prod-us-west-2/training-datasets/Amazon_SageMaker_FAQs/Amazon_SageMaker_FAQs.csv Amazon_SageMaker_FAQs.csv

import pandas as pd

data = pd.read_csv("Amazon_SageMaker_FAQs.csv", names=["Questions", "Answers"])
data["id"] = data.index
data_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}

Para hiperparámetros de entrenamiento específicos de algoritmos, el SDK de SageMaker se puede recuperar o sobrescribir:

from sagemaker import hyperparameters

hyperparameters = 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'}

La capacitación de SageMaker consta de dos pasos: crear el objeto estimador e iniciar el trabajo de capacitación. El resultado es un modelo empaquetado previamente con incorporaciones de su gran conjunto de datos utilizados como datos de entrenamiento, que se pueden implementar para inferencia y obtener el vecino más cercano para cualquier oración de entrada. Vea el siguiente código:

from sagemaker.jumpstart.estimator import JumpStartEstimator

estimator = 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()

La sintaxis de consulta para convertir texto en incrustaciones es la misma que antes. Sin embargo, el código para obtener el vecino más cercano se puede simplificar de la siguiente manera:

payload_nearest_neighbour = {
    "queries": ["Is R supported with 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"])]
>>> "Yes, R is supported with Amazon SageMaker. You can use R within SageMaker notebook instances, which include a preinstalled R kernel and the reticulate library. Reticulate offers an R interface for the Amazon SageMaker Python SDK, enabling ML practitioners to build, train, tune, and deploy R models."

También podemos consultar el punto final con preguntas en el Amazon_SageMaker_FAQs conjunto de datos y comparar cuántas de las respuestas correctas correspondientes se devuelven. En el siguiente ejemplo, medimos la precisión de los 3 primeros, dado que podría haber pares de preguntas y respuestas similares. Esto significa que si la respuesta correcta se devuelve como una de las 3 respuestas principales, se trata como una consulta correcta.

total_correct_answers = 0

for 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

Ejecute una transformación por lotes para obtener incrustaciones en grandes conjuntos de datos

Para empresas y organizaciones con un gran volumen de documentos históricos que exceden la memoria de una única instancia de punto final, puede utilizar la transformación por lotes de SageMaker para ahorrar costos. Cuando inicia un trabajo de transformación por lotes, SageMaker inicia los recursos informáticos necesarios para procesar los datos. Durante el trabajo, SageMaker aprovisiona y administra automáticamente los recursos informáticos. Cuando se completa el trabajo de transformación por lotes, esos recursos se limpian automáticamente, lo que minimiza los costos. Al dividir un conjunto de datos grande en fragmentos más pequeños y utilizar más instancias, puede escalar la computación para lograr una inferencia más rápida con un costo similar, sin administrar la infraestructura. La carga útil máxima para la transformación por lotes es de 100 MB y el tiempo de espera es de 1 hora.

El formato de entrada para nuestro trabajo de transformación por lotes es un archivo JSONL, con entradas como una línea de JSON, que consta de id y text_inputs. Vea el siguiente código:

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

Cuando los datos estén listos en Servicio de almacenamiento simple de Amazon (Amazon S3), puede crear el objeto de transformación por lotes desde el modelo JumpStart de SageMaker, que activa el trabajo de transformación:

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

Una vez completado el trabajo de transformación por lotes, puede descargar el resultado desde 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)

Conclusión

SageMaker JumpStart proporciona una forma sencilla de utilizar modelos básicos de lenguaje de última generación para la incrustación de texto y la búsqueda semántica. Con la interfaz de usuario o solo unas pocas líneas de código, puede implementar un modelo de incrustación de texto de alta precisión y encontrar coincidencias semánticas en grandes conjuntos de datos, a escala y de manera rentable. SageMaker JumpStart elimina las barreras para implementar la búsqueda semántica al brindar acceso instantáneo a modelos de vanguardia como los evaluados en la tabla de clasificación MTEB. Las empresas y los desarrolladores pueden crear sistemas inteligentes de búsqueda y recomendación más rápidamente.

Esta publicación demostró cómo encontrar preguntas y respuestas semánticamente similares, que podrían aplicarse a casos de uso, recomendaciones y personalización de RAG, traducciones multilingües y más. Con avances continuos en los modelos de lenguaje y la simplicidad de SageMaker JumpStart, más organizaciones pueden incorporar capacidades de IA generativa en sus productos. Como siguiente paso, puede probar modelos de incrustación de texto de SageMaker JumpStart en su propio conjunto de datos para probar y comparar los resultados para sus casos de uso de RAG.

Sobre los autores

Dr. Baichuan Sun, que actualmente se desempeña como arquitecto senior de soluciones de IA/ML en AWS, se enfoca en IA generativa y aplica su conocimiento en ciencia de datos y aprendizaje automático para brindar soluciones comerciales prácticas basadas en la nube. Con experiencia en consultoría de gestión y arquitectura de soluciones de inteligencia artificial, aborda una variedad de desafíos complejos, incluida la visión por computadora robótica, el pronóstico de series temporales y el mantenimiento predictivo, entre otros. Su trabajo se basa en una sólida formación en gestión de proyectos, investigación y desarrollo de software y actividades académicas. Fuera del trabajo, el Dr. Sun disfruta del equilibrio entre viajar y pasar tiempo con familiares y amigos, lo que refleja un compromiso tanto con su crecimiento profesional como con su bienestar personal.

Hemant Singh es un Científico Aplicado con experiencia en Amazon SageMaker JumpStart. Obtuvo su maestría en el Instituto Courant de Ciencias Matemáticas y su B.Tech en el IIT Delhi. Tiene experiencia trabajando en una amplia gama de problemas de aprendizaje automático dentro del dominio del procesamiento del lenguaje natural, la visión por computadora y el análisis de series temporales.

Dr. Ashish Khetan es un científico aplicado senior con algoritmos integrados de Amazon SageMaker y ayuda a desarrollar algoritmos de aprendizaje automático. Obtuvo su doctorado en la Universidad de Illinois Urbana-Champaign. Es un investigador activo en aprendizaje automático e inferencia estadística, y ha publicado numerosos artículos en conferencias NeurIPS, ICML, ICLR, JMLR, ACL y EMNLP.