Photo de real-napster sur Pixabay

Dans l'un de mes projets récents, j'ai dû créer un système de recherche sémantique capable d'évoluer avec des performances élevées et de fournir des réponses en temps réel pour les recherches de rapports. Nous avons utilisé PostgreSQL avec pgvector sur AWS RDS, associé à AWS Lambda, pour y parvenir. Le défi était de permettre aux utilisateurs d'effectuer des recherches à l'aide de requêtes en langage naturel au lieu de s'appuyer sur des mots-clés rigides, tout en garantissant que les réponses duraient moins de 1 à 2 secondes, voire moins, et ne pouvaient exploiter que les ressources du processeur.

Dans cet article, je passerai en revue les étapes que j'ai suivies pour créer ce système de recherche, de la récupération au reclassement, ainsi que les optimisations réalisées à l'aide d'OpenVINO et du traitement par lots intelligent pour la tokenisation.

Présentation de la recherche sémantique : récupération et reclassement

Les systèmes de recherche modernes de pointe se composent généralement de deux étapes principales : récupération et reclassement.

1) Récupération : La première étape consiste à récupérer un sous-ensemble de documents pertinents en fonction de la requête de l'utilisateur. Cela peut être fait à l'aide de modèles d'intégration pré-entraînés, tels que les intégrations petites et grandes d'OpenAI, les modèles Embed de Cohere ou les intégrations mxbai de Mixbread. La récupération se concentre sur la réduction du pool de documents en mesurant leur similarité avec la requête.

Voici un exemple simplifié utilisant la bibliothèque de transformation de phrases de Huggingface pour la récupération, qui est l'une de mes bibliothèques préférées pour cela :

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) Reclassement : Une fois les documents les plus pertinents récupérés, nous améliorons encore le classement de ces documents à l'aide d'un modèle de cross-encoder. Cette étape réévalue chaque document par rapport à la requête avec plus de précision, en se concentrant sur une compréhension contextuelle plus approfondie.
Le reclassement est bénéfique car il ajoute une couche de raffinement supplémentaire en notant plus précisément la pertinence de chaque document.

Voici un exemple de code pour le reclassement à l'aide de cross-encoder/ms-marco-TinyBERT-L-2-v2, un encodeur croisé léger :

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

Identifier les goulots d'étranglement : le coût de la tokenisation et de la prédiction

Au cours du développement, j'ai constaté que les étapes de tokenisation et de prédiction prenaient assez de temps lors de la gestion de 1 000 rapports avec des paramètres par défaut pour les transformateurs de phrases. Cela a créé un goulot d'étranglement en termes de performances, d'autant plus que nous visions des réponses en temps réel.

Ci-dessous j'ai profilé mon code à l'aide de SnakeViz pour visualiser les performances :

Comme vous pouvez le constater, les étapes de tokenisation et de prédiction sont disproportionnellement lentes, ce qui entraîne des retards importants dans la diffusion des résultats de recherche. Dans l’ensemble, cela a pris en moyenne 4 à 5 secondes. Cela est dû au fait qu’il existe des opérations de blocage entre les étapes de tokenisation et de prédiction. Si nous ajoutons également d'autres opérations comme l'appel à la base de données, le filtrage, etc., nous nous retrouvons facilement avec 8 à 9 secondes au total.

Optimiser les performances avec OpenVINO

La question à laquelle j'ai été confronté était : Pouvons-nous le rendre plus rapide ? La réponse est oui, en tirant parti de OpenVINO, un backend optimisé pour l'inférence CPU. OpenVINO permet d'accélérer l'inférence de modèles d'apprentissage profond sur le matériel Intel, que nous utilisons sur AWS Lambda.

Exemple de code pour l'optimisation OpenVINO
Voici comment j'ai intégré OpenVINO dans le système de recherche pour accélérer l'inférence :

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)

Avec cette approche, nous pourrions obtenir une accélération de 2 à 3 fois, réduisant les 4 à 5 secondes d'origine à 1 à 2 secondes. Le code de travail complet se trouve sur Github.

Optimisation de la vitesse : taille des lots et tokenisation

Un autre facteur essentiel pour améliorer les performances était l'optimisation du processus de tokénisation et l'ajustement de la taille du lot et de la longueur du jeton. En augmentant la taille du lot (batch_size=16) et en réduisant la longueur du jeton (max_length=512), nous pourrions paralléliser la tokenisation et réduire la surcharge des opérations répétitives. Dans nos expériences, nous avons constaté qu'une valeur batch_size comprise entre 16 et 64 fonctionnait bien, toute valeur supérieure dégradant les performances. De même, nous avons opté pour une max_length de 128, ce qui est viable si la longueur moyenne de vos rapports est relativement courte. Grâce à ces changements, nous avons atteint une accélération globale de 8x, réduisant le temps de reclassement à moins d'une seconde, même sur le processeur.

En pratique, cela signifiait expérimenter différentes tailles de lots et longueurs de jetons pour trouver le bon équilibre entre vitesse et précision pour vos données. Ce faisant, nous avons constaté des améliorations significatives des temps de réponse, rendant le système de recherche évolutif même avec 1 000 rapports.

Conclusion

En utilisant OpenVINO et en optimisant la tokenisation et le traitement par lots, nous avons pu créer un système de recherche sémantique hautes performances qui répond aux exigences en temps réel sur une configuration uniquement CPU. En fait, nous avons connu une accélération globale de 8x. La combinaison de la récupération à l'aide de transformateurs de phrases et du reclassement avec un modèle d'encodeur croisé crée une expérience de recherche puissante et conviviale.

Si vous construisez des systèmes similaires avec des contraintes de temps de réponse et de ressources de calcul, je vous recommande fortement d'explorer OpenVINO et le traitement par lots intelligent pour débloquer de meilleures performances.

J'espère que vous avez apprécié cet article. Si vous avez trouvé cet article utile, donnez-moi un like pour que d'autres puissent le trouver également et partagez-le avec vos amis. Suivez-moi sur Linkedin pour rester au courant de mon travail. Merci d'avoir lu !