RAG mínimo con embeddings

En este notebook vamos a implementar el núcleo de un sistema de Retrieval-Augmented Generation (RAG): dividir documentos en fragmentos, construir embeddings, recuperar los fragmentos más relevantes para una pregunta y preparar un contexto que luego podría enviarse a un LLM.

Preparación del ambiente

Instalemos las librerías necesarias desde el archivo de requerimientos asociado.

[ ]:

!wget -q https://raw.githubusercontent.com/santiagxf/M72109/master/docs/document-understanding/rag-minimo.txt
%pip install -r rag-minimo.txt --quiet

Documentos de ejemplo

Para enfocarnos en el patrón, usaremos documentos pequeños escritos como texto. En un caso real, estos fragmentos podrían venir de PDFs, páginas HTML, OCR o una base documental interna.

[ ]:

documentos = [
    "El OCR permite convertir imágenes de documentos en texto editable. Es útil para facturas escaneadas, formularios y contratos en papel.",
    "LayoutLM combina texto y coordenadas espaciales para entender la estructura visual de un documento.",
    "RAG recupera fragmentos relevantes de una colección documental y los usa como contexto para responder preguntas.",
    "Los embeddings representan textos como vectores densos. Textos semánticamente parecidos tienden a quedar cerca en el espacio vectorial.",
    "El chunking divide documentos largos en fragmentos más pequeños para mejorar la recuperación y controlar el tamaño del prompt.",
]

Construyendo embeddings

Usaremos un modelo pequeño de Sentence Transformers. El modelo transforma cada fragmento en un vector numérico que luego podemos comparar con el vector de la pregunta.

[ ]:

from sentence_transformers import SentenceTransformer

modelo_embeddings = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
embeddings_documentos = modelo_embeddings.encode(documentos, normalize_embeddings=True)

embeddings_documentos.shape

La primera dimensión corresponde a la cantidad de fragmentos. La segunda corresponde al tamaño del embedding producido por el modelo.

Recuperando fragmentos relevantes

Ahora codificamos la pregunta y calculamos similitud coseno contra todos los fragmentos. Como normalizamos los embeddings, el producto punto funciona como similitud coseno.

[ ]:

import numpy as np
import pandas as pd

pregunta = "¿Cómo ayuda RAG a responder preguntas sobre documentos?"
embedding_pregunta = modelo_embeddings.encode([pregunta], normalize_embeddings=True)

similitudes = np.dot(embeddings_documentos, embedding_pregunta.T).ravel()
ranking = np.argsort(similitudes)[::-1]

resultados = pd.DataFrame({
    "fragmento": [documentos[i] for i in ranking],
    "similitud": [similitudes[i] for i in ranking],
})

resultados.head(3)

Construyendo el contexto

En un sistema RAG completo, los fragmentos recuperados se agregan al prompt del modelo generativo. Aquí construiremos ese prompt de forma explícita para ver qué información recibe el LLM.

[ ]:

top_k = 2
fragmentos_recuperados = [documentos[i] for i in ranking[:top_k]]

contexto = "\n\n".join(f"Fragmento {idx + 1}: {texto}" for idx, texto in enumerate(fragmentos_recuperados))

prompt = f"""
Responda la pregunta usando solamente el contexto.

Contexto:
{contexto}

Pregunta: {pregunta}
Respuesta:
"""

print(prompt)

Note que todavía no llamamos a ningún LLM. Este notebook se concentra en el paso de retrieval, que suele ser el componente más importante para que la respuesta final esté anclada en documentos reales.

Cierre

El ejemplo anterior implementa el esqueleto de RAG con pocas líneas: embeddings, similitud, recuperación y armado de contexto. Para documentos reales deberíamos agregar carga de archivos, chunking con solapamiento, metadatos, evaluación y, cuando la colección crece, un índice vectorial especializado.