Reclasificación con modelos de incrustación

Un código en Python para el reordenamiento de RAG

Índice

Reranking es el segundo paso en Generación Aumentada de Recuperación (RAG) sistemas, justo entre Recuperar y Generar.

Cubos eléctricos en espacio digital

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

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 descargamos los documentos de internet primero, la respuesta del sistema de búsqueda podría verse afectada por preferencias/algoritmos del proveedor de búsqueda, contenido patrocinado, optimización SEO, etc., por lo que necesitamos un reranking post-búsqueda.

Lo que estaba haciendo:

  • obteniendo embeddings para la consulta de búsqueda
  • obteniendo embeddings para cada documento. el documento no se esperaba que fuera más de 8k tokens de todos modos
  • calculando la similitud entre la consulta y cada uno de los embeddings del documento
  • ordenando los documentos por esta similitud.

No hay base de datos vectorial aquí, vamos.

Código de ejemplo

Usando 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 modelos LLM de embeddings el umbral será diferente.

Me alegraré si este fragmento de código es útil para ti de alguna manera. Licencia Copy/Paste/UseAnyWayYouWant. 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 embeddings

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

El segundo lugar lo obtuvo 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