Reranking med inbäddningsmodeller

En Python-kod för RAG:s omrankning

Sidinnehåll

Reranking är ett andra steg i Retrieval Augmented Generation (RAG) system, precis mellan Retrieval och Generating.

Elektriska kuber i digitalt utrymme

Detta ovan är hur Flux-1 dev föreställer sig Elektriska kuber i digitalt utrymme.

Retrieval med reranking

Om vi lagrar dokumenten i form av embeddingar i vektordatabasen från början - kommer retrievalen att ge oss en lista med liknande dokument direkt.

Fristående reranking

Men om vi laddar ner dokumenten från internet först, kan sökningssystemets svar påverkas av sökleverantörers preferenser/algoritmer, sponsrat innehåll, SEO-optimering etc., så vi behöver post-sökning reranking.

Vad jag gjorde -

  • fick embeddingar för sökfrågan
  • fick embeddingar för varje dokument. dokumentet förväntades inte vara mer än 8k tokens ändå
  • beräknade likheten mellan frågan och varje dokumentets embeddingar
  • sorterade dokumenten efter denna likhet.

Ingen vektordatabas här, låt oss gå.

Exempelkod

Använder Langchain för att ansluta till Ollama och Langchains cosine_similarity funktion. Du kan filtrera efter likhetsmått, men kom ihåg att för olika domäner och embedding LLMs kommer tröskeln att vara olika.

Jag skulle bli glad om denna bit kod är till nytta för dig på något sätt. Copy/Paste/UseAnyWayYouWant-licens. Hälsa.

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

Bäst Embedding Modeller

För mina uppgifter är den bästa embeddingmodellen för tillfället bge-large:335m-en-v1.5-fp16

Andra plats tog nomic-embed-text:137m-v1.5-fp16 och jina/jina-embeddings-v2-base-en:latest.

Men gör dina egna tester för din egen domän och frågor.

Användbara länkar