{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "RAG mínimo con embeddings\n",
    "=========================\n",
    "\n",
    "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."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Preparación del ambiente\n",
    "\n",
    "Instalemos las librerías necesarias desde el archivo de requerimientos asociado."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!wget -q https://raw.githubusercontent.com/santiagxf/M72109/master/docs/document-understanding/rag-minimo.txt\n",
    "%pip install -r rag-minimo.txt --quiet"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Documentos de ejemplo\n",
    "\n",
    "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."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "documentos = [\n",
    "    \"El OCR permite convertir imágenes de documentos en texto editable. Es útil para facturas escaneadas, formularios y contratos en papel.\",\n",
    "    \"LayoutLM combina texto y coordenadas espaciales para entender la estructura visual de un documento.\",\n",
    "    \"RAG recupera fragmentos relevantes de una colección documental y los usa como contexto para responder preguntas.\",\n",
    "    \"Los embeddings representan textos como vectores densos. Textos semánticamente parecidos tienden a quedar cerca en el espacio vectorial.\",\n",
    "    \"El chunking divide documentos largos en fragmentos más pequeños para mejorar la recuperación y controlar el tamaño del prompt.\",\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Construyendo embeddings\n",
    "\n",
    "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."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sentence_transformers import SentenceTransformer\n",
    "\n",
    "modelo_embeddings = SentenceTransformer(\"sentence-transformers/all-MiniLM-L6-v2\")\n",
    "embeddings_documentos = modelo_embeddings.encode(documentos, normalize_embeddings=True)\n",
    "\n",
    "embeddings_documentos.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "La primera dimensión corresponde a la cantidad de fragmentos. La segunda corresponde al tamaño del embedding producido por el modelo."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Recuperando fragmentos relevantes\n",
    "\n",
    "Ahora codificamos la pregunta y calculamos similitud coseno contra todos los fragmentos. Como normalizamos los embeddings, el producto punto funciona como similitud coseno."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "pregunta = \"¿Cómo ayuda RAG a responder preguntas sobre documentos?\"\n",
    "embedding_pregunta = modelo_embeddings.encode([pregunta], normalize_embeddings=True)\n",
    "\n",
    "similitudes = np.dot(embeddings_documentos, embedding_pregunta.T).ravel()\n",
    "ranking = np.argsort(similitudes)[::-1]\n",
    "\n",
    "resultados = pd.DataFrame({\n",
    "    \"fragmento\": [documentos[i] for i in ranking],\n",
    "    \"similitud\": [similitudes[i] for i in ranking],\n",
    "})\n",
    "\n",
    "resultados.head(3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Construyendo el contexto\n",
    "\n",
    "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."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "top_k = 2\n",
    "fragmentos_recuperados = [documentos[i] for i in ranking[:top_k]]\n",
    "\n",
    "contexto = \"\\n\\n\".join(f\"Fragmento {idx + 1}: {texto}\" for idx, texto in enumerate(fragmentos_recuperados))\n",
    "\n",
    "prompt = f\"\"\"\n",
    "Responda la pregunta usando solamente el contexto.\n",
    "\n",
    "Contexto:\n",
    "{contexto}\n",
    "\n",
    "Pregunta: {pregunta}\n",
    "Respuesta:\n",
    "\"\"\"\n",
    "\n",
    "print(prompt)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "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."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Cierre\n",
    "\n",
    "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."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}