Reclassificação com modelos de embedding

Um código Python para o reranking de RAG

Conteúdo da página

Reranking é uma segunda etapa na Geração Aumentada por Recuperação (RAG) sistemas, situada exatamente entre a Recuperação e a Geração.

Cubos elétricos no espaço digital

O acima é como o Flux-1 dev imagina Cubos elétricos no espaço digital.

Para um guia completo sobre a construção de sistemas RAG, consulte o Tutorial de Geração Aumentada por Recuperação (RAG): Arquitetura, Implementação e Guia de Produção.

Este artigo aborda o reranking com modelos de embedding. Para abordagens alternativas, você também pode conferir Reranking de textos com Ollama e Qwen3 Embedding LLM - em Go e Reranking de documentos com Ollama e modelo Reranker Qwen3 - em Go.

Recuperação com reranking

Se armazenarmos os documentos na forma de embeddings no banco de dados vetorial desde o início, a recuperação nos dará a lista de documentos similares imediatamente.

Reranking independente

Mas se baixarmos os documentos da internet primeiro, a resposta do sistema de busca pode ser afetada por preferências/algoritmos do provedor de busca, conteúdo patrocinado, otimização de SEO, etc., então precisamos fazer o reranking pós-busca.

O que eu estava fazendo:

  • obter embeddings para a consulta de busca
  • obter embeddings para cada documento. o documento não era esperado para ter mais de 8k tokens de qualquer maneira
  • calcular a similaridade entre a consulta e os embeddings de cada documento
  • ordenar os documentos por esta similaridade.

Sem banco de dados vetorial aqui, vamos lá.

Código de exemplo

Usando o Langchain para se conectar ao Ollama e a função cosine_similarity do langchain. Você pode filtrar pela medida de similaridade, mas tenha em mente que para diferentes domínios e LLMs de embedding, o limite seria diferente.

Ficarei feliz se este pedaço de código for útil para você de alguma forma. Licença Copiar/Colar/UsarComoQuiser. Saúde.

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

Melhores Modelos de Embedding

Para as minhas tarefas, o melhor modelo de embedding atualmente é bge-large:335m-en-v1.5-fp16.

O segundo lugar foi para nomic-embed-text:137m-v1.5-fp16 e jina/jina-embeddings-v2-base-en:latest.

Mas faça seus próprios testes para o seu próprio domínio e consultas.