Reranking avec des modèles d'embedding

Un code Python pour le reranking de RAG

Sommaire

Reranking est une deuxième étape dans la Génération Augmentée par Récupération (RAG) systèmes, juste entre la Récupération et la Génération.

Cubes électriques dans l’espace numérique

Ce qui est représenté ci-dessus est la manière dont Flux-1 dev imagine Cubes électriques dans l'espace numérique.

Récupération avec reranking

Si nous stockons les documents sous forme d’embeddings dans la base de données vectorielle dès le début – la récupération nous donnera directement la liste des documents similaires.

Reranking autonome

Mais si nous téléchargeons d’abord les documents depuis l’internet, la réponse du système de recherche pourrait être influencée par les préférences/algorithmes des fournisseurs de recherche, le contenu sponsorisé, l’optimisation SEO, etc. Nous aurons donc besoin d’un reranking post-recherche.

Ce que j’ai fait :

  • obtenir les embeddings pour la requête de recherche
  • obtenir les embeddings pour chaque document. Le document n’était pas censé dépasser 8 000 tokens de toute façon
  • calculer la similarité entre la requête et chaque embedding du document
  • trier les documents selon cette similarité.

Aucune base de données vectorielle ici, allons-y.

Exemple de code

Utilisation de Langchain pour se connecter à Ollama et à la fonction cosine_similarity de langchain. Vous pouvez filtrer par mesure de similarité, mais gardez à l’esprit que pour différents domaines et modèles LLM d’embedding, le seuil sera différent.

Je serai heureux si ce bout de code vous est utile de quelque manière que ce soit. Licence Copy/Paste/UseAnyWayYouWant. Merci.

from langchain_core.documents import Document
from langchain_community.embeddings import OllamaEmbeddings
from langchain_community.utils.math import cosine_similarity
import numpy as np


def cosine_distance(a: np.ndarray, b: np.ndarray) -> np.ndarray:
    return 1.0 - cosine_similarity(a, b)

def compute_score(vectors: np.ndarray) -> float:
    score = cosine_distance(vectors[0].reshape(1, -1), vectors[1].reshape(1, -1)).item()
    return score

def list_to_array(lst):
    return np.array(lst, dtype=float)   

def compute_scorel(lists) -> float:
    v1 = list_to_array(lists[0])
    v2 = list_to_array(lists[1])
    return compute_score([v1, v2])

def filter_docs(emb_model_name, docs, query, num_docs):
    content_arr = [doc.page_content for doc in docs]

    ollama_emb = OllamaEmbeddings(
        model=emb_model_name
    )

    docs_embs = ollama_emb.embed_documents(content_arr)
    query_embs = ollama_emb.embed_query(query)
    sims = []
    for i, emb in enumerate(docs_embs):
        idx = docs[i].id
        s = compute_scorel([query_embs, docs_embs[i]])
        simstr = str(round(s, 4))
        docs[i].metadata["sim"] = simstr
        sim = {
            "idx": idx,
            "i": i,
            "sim": s,
        }
        sims.append(sim)

    sims.sort(key=sortFn)

    sorted_docs = [docs[x["i"]] for x in sims]
    filtered_docs = sorted_docs[:num_docs]
    return filtered_docs

Meilleurs modèles d’embedding

Pour mes tâches, le meilleur modèle d’embedding actuellement est bge-large:335m-en-v1.5-fp16

La deuxième place revient à nomic-embed-text:137m-v1.5-fp16 et jina/jina-embeddings-v2-base-en:latest.

Mais faites vos propres tests pour votre propre domaine et vos propres requêtes.

Liens utiles