Reranking mit Embedding-Modellen

Ein Python-Code für das Wiederholungsrangieren von RAG

Inhaltsverzeichnis

Reranking ist der zweite Schritt im Retrieval Augmented Generation (RAG) systemen, direkt zwischen Retrieving und Generating.

Elektrische Würfel im digitalen Raum

Das oben ist, wie Flux-1 dev sich Elektrische Würfel im digitalen Raum vorstellt.

Retrieval mit Reranking

Wenn wir die Dokumente von Anfang an in Form von Embeddings in der Vektor-Datenbank speichern – wird der Retrieval uns direkt die Liste ähnlicher Dokumente liefern.

Standalone Reranking

Aber wenn wir die Dokumente zuerst aus dem Internet herunterladen, könnte die Antwort des Suchsystems von den Vorlieben/Algorithmen des Suchanbieters, sponsored Content, SEO-Optimierung usw. beeinflusst werden. Daher benötigen wir ein Post-Search-Reranking.

Was ich tat:

  • Embeddings für die Suchanfrage erhalte
  • Embeddings für jedes Dokument. Das Dokument wurde ohnehin nicht erwartet, länger als 8k Tokens zu sein
  • Ähnlichkeit zwischen Anfrage und jedem Dokument-Embedding berechnen
  • Dokumente nach dieser Ähnlichkeit sortieren

Keine Vektor-Datenbank hier, los geht’s.

Beispielcode

Mit Langchain, um mit Ollama zu verbinden und die cosine_similarity-Funktion von langchain zu verwenden. Sie können nach Ähnlichkeitsmaß filtern, beachten Sie jedoch, dass der Schwellenwert für verschiedene Domänen und Embedding-LLMs unterschiedlich sein wird.

Ich wäre froh, wenn dieser Codeabschnitt Ihnen auf irgendeine Weise nützlich ist. Lizenz: Copy/Paste/UseAnyWayYouWant. Viel Spaß.

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

Beste Embedding-Modelle

Für meine Aufgaben ist derzeit das beste Embedding-Modell bge-large:335m-en-v1.5-fp16

Auf dem zweiten Platz folgen nomic-embed-text:137m-v1.5-fp16 und jina/jina-embeddings-v2-base-en:latest.

Aber führen Sie Ihre eigenen Tests für Ihre eigenen Domänen und Abfragen durch.