Foto de real-napster en Pixabay

En uno de mis proyectos recientes, tuve que crear un sistema de búsqueda semántica que pudiera escalar con un alto rendimiento y ofrecer respuestas en tiempo real para búsquedas de informes. Usamos PostgreSQL con pgvector en AWS RDS, junto con AWS Lambda, para lograr esto. El desafío era permitir a los usuarios realizar búsquedas mediante consultas en lenguaje natural en lugar de depender de palabras clave rígidas, y al mismo tiempo garantizar que las respuestas duraran menos de 1 o 2 segundos o incluso menos y solo pudieran aprovechar los recursos de la CPU.

En esta publicación, explicaré los pasos que seguí para crear este sistema de búsqueda, desde la recuperación hasta la reclasificación, y las optimizaciones realizadas con OpenVINO y el procesamiento por lotes inteligente para la tokenización.

Descripción general de la búsqueda semántica: recuperación y reclasificación

Los sistemas de búsqueda modernos de última generación generalmente constan de dos pasos principales: recuperación y reclasificación.

1) Recuperación: El primer paso implica recuperar un subconjunto de documentos relevantes según la consulta del usuario. Esto se puede hacer utilizando modelos de incrustaciones previamente entrenados, como las incrustaciones pequeñas y grandes de OpenAI, los modelos Embed de Cohere o las incrustaciones mxbai de Mixbread. La recuperación se centra en reducir el conjunto de documentos midiendo su similitud con la consulta.

Aquí hay un ejemplo simplificado que utiliza la biblioteca de transformadores de oraciones de Huggingface para la recuperación, que es una de mis bibliotecas favoritas para esto:

from sentence_transformers import SentenceTransformer
import numpy as np

# Load a pre-trained sentence transformer model
model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")

# Sample query and documents (vectorize the query and the documents)
query = "How do I fix a broken landing gear?"
documents = ["Report 1 on landing gear failure", "Report 2 on engine problems"]

# Get embeddings for query and documents
query_embedding = model.encode(query)
document_embeddings = model.encode(documents)

# Calculate cosine similarity between query and documents
similarities = np.dot(document_embeddings, query_embedding)

# Retrieve top-k most relevant documents
top_k = np.argsort(similarities)[-5:]
print("Top 5 documents:", [documents[i] for i in top_k])

2) Reclasificación: Una vez que se han recuperado los documentos más relevantes, mejoramos aún más la clasificación de estos documentos utilizando un modelo codificador cruzado. Este paso reevalúa cada documento en relación con la consulta con mayor precisión, enfocándose en una comprensión contextual más profunda.
Reclasificar es beneficioso porque agrega una capa adicional de refinamiento al calificar la relevancia de cada documento con mayor precisión.

Aquí hay un ejemplo de código para reclasificar usando cross-encoder/ms-marco-TinyBERT-L-2-v2, un codificador cruzado liviano:

from sentence_transformers import CrossEncoder

# Load the cross-encoder model
cross_encoder = CrossEncoder("cross-encoder/ms-marco-TinyBERT-L-2-v2")

# Use the cross-encoder to rerank top-k retrieved documents
query_document_pairs = [(query, doc) for doc in documents]
scores = cross_encoder.predict(query_document_pairs)

# Rank documents based on the new scores
top_k_reranked = np.argsort(scores)[-5:]
print("Top 5 reranked documents:", [documents[i] for i in top_k_reranked])

Identificación de cuellos de botella: el costo de la tokenización y la predicción

Durante el desarrollo, descubrí que las etapas de tokenización y predicción estaban tardando bastante al manejar 1000 informes con configuraciones predeterminadas para transformadores de oraciones. Esto creó un cuello de botella en el rendimiento, especialmente porque buscábamos respuestas en tiempo real.

A continuación, perfilé mi código usando SnakeViz para visualizar las actuaciones:

Como puede ver, los pasos de tokenización y predicción son desproporcionadamente lentos, lo que genera retrasos significativos en la entrega de resultados de búsqueda. En general, tomó entre 4 y 5 segundos en promedio. Esto se debe al hecho de que existen operaciones de bloqueo entre los pasos de tokenización y predicción. Si también sumamos otras operaciones como llamadas a bases de datos, filtrado, etc., fácilmente terminamos con 8-9 segundos en total.

Optimización del rendimiento con OpenVINO

La pregunta que enfrenté fue: ¿Podemos hacerlo más rápido? La respuesta es sí, aprovechando OpenVINO, un backend optimizado para la inferencia de CPU. OpenVINO ayuda a acelerar la inferencia de modelos de aprendizaje profundo en hardware Intel, que utilizamos en AWS Lambda.

Ejemplo de código para optimización de OpenVINO
Así es como integré OpenVINO en el sistema de búsqueda para acelerar la inferencia:

import argparse
import numpy as np
import pandas as pd
from typing import Any
from openvino.runtime import Core
from transformers import AutoTokenizer


def load_openvino_model(model_path: str) -> Core:
    core = Core()
    model = core.read_model(model_path   ".xml")
    compiled_model = core.compile_model(model, "CPU")
    return compiled_model


def rerank(
    compiled_model: Core,
    query: str,
    results: list[str],
    tokenizer: AutoTokenizer,
    batch_size: int,
) -> np.ndarray[np.float32, Any]:
    max_length = 512
    all_logits = []

    # Split results into batches
    for i in range(0, len(results), batch_size):
        batch_results = results[i : i   batch_size]
        inputs = tokenizer(
            [(query, item) for item in batch_results],
            padding=True,
            truncation="longest_first",
            max_length=max_length,
            return_tensors="np",
        )

        # Extract input tensors (convert to NumPy arrays)
        input_ids = inputs["input_ids"].astype(np.int32)
        attention_mask = inputs["attention_mask"].astype(np.int32)
        token_type_ids = inputs.get("token_type_ids", np.zeros_like(input_ids)).astype(
            np.int32
        )

        infer_request = compiled_model.create_infer_request()
        output = infer_request.infer(
            {
                "input_ids": input_ids,
                "attention_mask": attention_mask,
                "token_type_ids": token_type_ids,
            }
        )

        logits = output["logits"]
        all_logits.append(logits)

    all_logits = np.concatenate(all_logits, axis=0)
    return all_logits


def fetch_search_data(search_text: str) -> pd.DataFrame:
    # Usually you would fetch the data from a database
    df = pd.read_csv("cnbc_headlines.csv")
    df = df[~df["Headlines"].isnull()]

    texts = df["Headlines"].tolist()

    # Load the model and rerank
    openvino_model = load_openvino_model("cross-encoder-openvino-model/model")
    tokenizer = AutoTokenizer.from_pretrained("cross-encoder/ms-marco-TinyBERT-L-2-v2")
    rerank_scores = rerank(openvino_model, search_text, texts, tokenizer, batch_size=16)

    # Add the rerank scores to the DataFrame and sort by the new scores
    df["rerank_score"] = rerank_scores
    df = df.sort_values(by="rerank_score", ascending=False)

    return df


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Fetch search results with reranking using OpenVINO"
    )

    parser.add_argument(
        "--search_text",
        type=str,
        required=True,
        help="The search text to use for reranking",
    )

    args = parser.parse_args()

    df = fetch_search_data(args.search_text)
    print(df)

Con este enfoque podríamos obtener una aceleración de 2 a 3 veces, reduciendo los 4 a 5 segundos originales a 1 a 2 segundos. El código de trabajo completo está en Github.

Ajuste de velocidad: tamaño de lote y tokenización

Otro factor crítico para mejorar el rendimiento fue optimizar el proceso de tokenización y ajustar el tamaño del lote y la longitud del token. Al aumentar el tamaño del lote (batch_size=16) y reducir la longitud del token (max_length=512), podríamos paralelizar la tokenización y reducir la sobrecarga de las operaciones repetitivas. En nuestros experimentos, descubrimos que un tamaño de lote entre 16 y 64 funcionaba bien, y un valor mayor degradaba el rendimiento. De manera similar, nos decidimos por una longitud máxima de 128, que es viable si la longitud promedio de sus informes es relativamente corta. Con estos cambios, logramos una aceleración general de 8 veces, reduciendo el tiempo de reclasificación a menos de 1 segundo, incluso en la CPU.

En la práctica, esto significó experimentar con diferentes tamaños de lotes y longitudes de tokens para encontrar el equilibrio adecuado entre velocidad y precisión para sus datos. Al hacerlo, vimos mejoras significativas en los tiempos de respuesta, lo que hizo que el sistema de búsqueda fuera escalable incluso con 1000 informes.

Conclusión

Al utilizar OpenVINO y optimizar la tokenización y el procesamiento por lotes, pudimos crear un sistema de búsqueda semántica de alto rendimiento que cumple con los requisitos en tiempo real en una configuración de solo CPU. De hecho, experimentamos una aceleración general de 8 veces. La combinación de recuperación mediante transformadores de oraciones y reclasificación con un modelo de codificador cruzado crea una experiencia de búsqueda poderosa y fácil de usar.

Si está creando sistemas similares con limitaciones en el tiempo de respuesta y los recursos computacionales, le recomiendo explorar OpenVINO y el procesamiento por lotes inteligente para desbloquear un mejor rendimiento.

Esperamos que hayas disfrutado este artículo. Si este artículo te resultó útil, dame un me gusta para que otros también puedan encontrarlo y compártelo con tus amigos. Sígueme en Linkedin para estar al día de mi trabajo. ¡Gracias por leer!