Sci Math && Code
Blog / Cómo construir un sistema RAG localmente
13 de abril de 2024
nlpllmragollama

Cómo construir un sistema RAG localmente

Construye un sistema RAG localmente usando Ollama y PostgreSQL con PgVector

Cómo construir un sistema RAG localmente

¿Qué es un sistema RAG?

Un sistema de Generación Aumentada por Recuperación (RAG, por sus siglas en inglés) combina las capacidades de un modelo de lenguaje de gran escala con un componente de recuperación que puede obtener documentos o pasajes relevantes de un corpus. Esta combinación permite que el modelo genere respuestas fluidas e informadas no solo basándose en su conocimiento entrenado, sino también recuperando y referenciando información factual de los documentos suministrados.

¿Cómo funciona?

Un sistema RAG se compone de dos partes principales: un motor de recuperación y un modelo de lenguaje de gran escala.

Cuando el usuario proporciona una consulta, el motor de recuperación busca en el corpus los pasajes o la información más relevantes. Típicamente esto se hace con técnicas de búsqueda semántica o similitud vectorial para rankear los documentos según su relevancia.

Los documentos mejor rankeados se formatean en una ventana de contexto que el modelo de lenguaje puede consumir. Esta ventana proporciona al modelo información relevante y hechos de los documentos recuperados.

Luego, el modelo toma la consulta original del usuario más la ventana de contexto como entrada. Combinando su propio conocimiento entrenado con la información suplementaria de los documentos recuperados, puede generar una respuesta fluida e informativa.

Finalmente, el sistema devuelve esa respuesta como la respuesta definitiva al usuario.

¿Por qué construirlo localmente?

Tener un sistema RAG corriendo localmente tiene varias ventajas. Primero, te permite experimentar dentro de tu propio entorno sin depender de servicios o APIs externos. Esto es especialmente útil para pruebas, debugging y personalización. Además, un sistema RAG local mejora la privacidad y seguridad de los datos, ya que la información sensible permanece en tu infraestructura. Por último, correr el sistema localmente puede ofrecer ventajas de rendimiento al reducir la latencia y eliminar la dependencia de condiciones de red externas.

¡Empecemos!

Para construir el sistema RAG localmente necesitamos dos componentes principales: un Modelo de Lenguaje de Gran Escala (LLM) y un motor de recuperación. El primer paso es configurar un LLM en nuestra máquina local.

El LLM

Estamos construyendo esto en un ordenador personal o estación de trabajo básica, por lo que necesitamos un LLM relativamente ligero en requisitos de recursos. Para lograrlo, utilizaremos Ollama.

Si bien Ollama puede aprovechar la aceleración GPU, no requiere GPUs especializadas para deep learning. Esto lo convierte en una elección ideal para nuestro sistema RAG local.

Ollama es una herramienta avanzada de IA que permite a los usuarios ejecutar modelos de lenguaje de gran escala (LLM) localmente en sus computadoras. Simplifica el proceso de ejecución local de modelos de lenguaje, proporcionando mayor control y flexibilidad. Soporta una variedad de modelos, incluyendo Llama 2, Mistral y otros.

Para comenzar, visita https://ollama.com y sigue las instrucciones de instalación.

PostgreSQL como base de datos vectorial

Para nuestras necesidades de almacenamiento vectorial, utilizaremos PostgreSQL junto con la extensión pgvector. PgVector es una extensión open-source para PostgreSQL que permite almacenar y buscar sobre embeddings generados por machine learning de forma eficiente.

¿Por qué una base de datos vectorial?

En el núcleo de un sistema RAG está la capacidad de recuperar rápidamente documentos relevantes a partir de una consulta. Esto se logra convirtiendo la consulta y los documentos en representaciones vectoriales de alta dimensión llamadas embeddings. Luego, la base de datos vectorial encuentra los embeddings más similares — los “vecinos más cercanos” — a la consulta, lo que nos permite recuperar los pasajes más relevantes.

Configurando PostgreSQL con pgvector

Crea un directorio nuevo para el proyecto. Luego crea un archivo docker-compose.yml:

version: "3.8"

services:
  db:
    image: pgvector/pgvector:pg16
    restart: always
    env_file:
      - .env
    ports:
      - ${POSTGRES_PORT}:${POSTGRES_PORT}
    environment:
      - PGDATA=/var/lib/postgresql/data/pgdata
      - POSTGRES_PORT=${POSTGRES_PORT}
      - POSTGRES_DB=${POSTGRES_DB}
      - POSTGRES_USER=${POSTGRES_USER}
      - POSTGRES_PASSWORD=${POSTGRES_PASSWORD}
    volumes:
      - pgdata:/var/lib/postgresql/data/pgdata

volumes:
  pgdata:

Y un archivo .env:

POSTGRES_SERVER=localhost
POSTGRES_PORT=5432
POSTGRES_DB=localrag
POSTGRES_USER=postgres
POSTGRES_PASSWORD=pgpassword

Levanta el contenedor:

docker-compose up -d

Instalando las librerías

Prefiero pipenv, pero puedes usar pip, poetry o cualquier otra herramienta. Crea un entorno con Python 3.12:

pipenv --python 3.12
pipenv install langchain langchain-community wikipedia python-dotenv pgvector psycopg2-binary

Crea un archivo main.py e importa las librerías necesarias:

import os
from langchain_community.document_loaders.wikipedia import WikipediaLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.embeddings.ollama import OllamaEmbeddings
from langchain_community.llms.ollama import Ollama
from langchain.vectorstores.pgvector import PGVector
from langchain.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough
from dotenv import load_dotenv

load_dotenv("../.env")

Define los modelos y la URI de base de datos:

EMBEDDING_MODEL = "mxbai-embed-large"
LLM_MODEL = "gemma:2b"

def database_uri():
    user = os.getenv("POSTGRES_USER", "postgres")
    password = os.getenv("POSTGRES_PASSWORD", "")
    server = os.getenv("POSTGRES_SERVER", "localhost")
    port = os.getenv("POSTGRES_PORT", "5432")
    db = os.getenv("POSTGRES_DB", "localrag")
    return f"postgresql+psycopg2://{user}:{password}@{server}:{port}/{db}"

Instanciando los componentes principales

embeddings = OllamaEmbeddings(model=EMBEDDING_MODEL)
vector_db = PGVector(
    embedding_function=embeddings,
    connection_string=database_uri(),
    pre_delete_collection=True,
)
retriever = vector_db.as_retriever()
llm = Ollama(model=LLM_MODEL)
  • Modelo de embeddings — genera representaciones vectoriales del texto usando mxbai-embed-large.
  • Base de datos vectorial — la instancia PGVector almacena y recupera embeddings desde PostgreSQL (pre_delete_collection=True es solo para desarrollo; nunca lo uses en producción).
  • Retriever — realiza búsquedas de similitud en la base de datos vectorial.
  • Modelo de lenguaje — la instancia de Ollama con gemma:2b para generar respuestas.

Encadenando los componentes RAG

human_prompt = """
Eres un asistente para tareas de respuesta a preguntas. Usa los siguientes fragmentos de contexto recuperado para responder la pregunta.
Si no sabes la respuesta, simplemente di que no lo sabes.
Usa un máximo de tres oraciones y mantén la respuesta concisa.\n
Contexto: {context} \n
Pregunta: {question} \n
Respuesta:
"""

prompt = ChatPromptTemplate.from_messages([("human", human_prompt)])

rag_chain = (
    {"context": retriever, "question": RunnablePassthrough()}
    | prompt
    | llm
    | StrOutputParser()
)

El pipeline funciona así:

  1. La pregunta del usuario se pasa al retriever, que obtiene los pasajes más relevantes de la base de datos vectorial.
  2. La pregunta y el contexto recuperado se inyectan en la plantilla de prompt.
  3. El prompt formateado se pasa al LLM, que genera una respuesta.
  4. StrOutputParser garantiza que la salida sea un string simple.

Probando el sistema RAG

def wikipedia_query(query: str):
    loader = WikipediaLoader(query=query, load_max_docs=3)
    docs = loader.load()
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=80)
    texts = text_splitter.split_documents(docs)
    vector_db.add_documents(texts)

# Consultas de prueba
wikipedia_query("Transformer Architecture")
result = rag_chain.invoke("¿En qué año se publicó el paper del Transformer?")
print(f"Consulta: ¿En qué año se publicó el paper del Transformer?\nResultado: {result}")

wikipedia_query("Langchain")
result = rag_chain.invoke("¿Qué es LangChain?")
print(f"Consulta: ¿Qué es LangChain?\nResultado: {result}")

Estos ejemplos usan documentos de Wikipedia para las pruebas. En un escenario real, reemplázalos con tus propios documentos específicos del dominio: PDFs, registros de chat, documentación interna, etc.

¡Eso es todo!

Si quieres explorar una implementación más completa, incluyendo una API funcional, visita mi repositorio de GitHub: https://github.com/luiscib3r/LLM-Projects/tree/main/local-rag. Encontrarás un Jupyter Notebook completo y una implementación de API lista para usar.