Reordenamiento con modelos de incrustación

Un código Python para el reordenamiento en RAG

Índice

Reranking es un segundo paso en la Generación Aumentada por Recuperación (RAG) sistemas, situado justo entre la Recuperación y la Generación.

Cubos eléctricos en el espacio digital

Lo anterior es cómo Flux-1 dev imagina Cubos eléctricos en el espacio digital.

Para una guía completa sobre la construcción de sistemas RAG, consulta el Tutorial de Generación Aumentada por Recuperación (RAG): Arquitectura, Implementación y Guía de Producción.

Este artículo cubre el reranking con modelos de embedding. Para enfoques alternativos, también puedes consultar Reranking de textos con Ollama y Qwen3 Embedding LLM - en Go y Reranking de documentos con Ollama y modelo Reranker Qwen3 - en Go.

Recuperación con reranking

Si almacenamos los documentos en forma de embeddings en la base de datos vectorial desde el principio, la recuperación nos dará la lista de documentos similares de inmediato.

Reranking independiente

Pero si primero descargamos los documentos de internet, la respuesta del sistema de búsqueda podría verse afectada por las preferencias/algoritmos del proveedor de búsqueda, contenido patrocinado, optimización SEO, etc., por lo que necesitamos realizar un reranking post-búsqueda.

Lo que estaba haciendo:

  • obtener embeddings para la consulta de búsqueda
  • obtener embeddings para cada documento. de todos modos, no se esperaba que el documento tuviera más de 8k tokens
  • calcular la similitud entre la consulta y los embeddings de cada documento
  • ordenar los documentos por esta similitud.

Sin base de datos vectorial aquí, ¡vamos allá!

Código de ejemplo

Utilizando Langchain para conectarse a Ollama y la función cosine_similarity de langchain. Puedes filtrar por medida de similitud, pero ten en cuenta que para diferentes dominios y LLM de embedding el umbral sería diferente.

Me gustaría que este fragmento de código te sea útil de alguna manera. Licencia Copiar/Pegar/UsarComoQuieras. Saludos.

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

Mejores Modelos de Embedding

Para mis tareas, el mejor modelo de embedding actualmente es bge-large:335m-en-v1.5-fp16.

El segundo lugar lo ocuparon nomic-embed-text:137m-v1.5-fp16 y jina/jina-embeddings-v2-base-en:latest.

Pero haz tus propias pruebas para tu propio dominio y consultas.

Enlaces útiles