Parent Document Retriever en Acción: Configuración de RAG con Mistral LLM y LangChain

Introducción

Este post demuestra cómo configurar un sistema de Generación Aumentada por Recuperación (RAG) utilizando LangChain, integrando un Parent Document Retriever con modelos de Mistral AI. Proporciona detalles de implementación, incluyendo código en Python, para mostrar cómo RAG mejora la calidad de las respuestas de los modelos de lenguaje. Esto complementa el post relacionado, que presenta el caso final.

Arquitectura Clave de RAG

Aquí están los componentes clave de este sistema RAG, describiendo sus roles y contribuciones a la arquitectura general:

Mistral LLM y embeddings: Utiliza embeddings de Mistral y el modelo Mistral Large LLM para generar respuestas relevantes basadas en conocimiento externo.
Parent Document Retriever: Recupera fragmentos más pequeños de información mientras hace referencia a sus documentos principales para obtener contexto.
FAISS Vector Store: Almacena embeddings y permite búsquedas de similitud eficientes.
Document Chunking: Divide documentos PDF en partes más pequeñas para una mejor recuperación.
Naive RAG Chain: Conecta el retriever, el vector store y el LLM para generar respuestas informadas.

Con estos componentes en mente, exploremos ahora cómo trabajan juntos para crear un sistema efectivo de Generación Aumentada por Recuperación.

Framework LangChain

LangChain es el framework principal utilizado para implementar el sistema de Generación Aumentada por Recuperación (RAG). Proporciona las herramientas y componentes necesarios para construir la aplicación RAG, incluyendo la integración con vector stores y retrievers.

Ahora que tenemos un framework, profundicemos en cómo el ParentDocumentRetriever equilibra la especificidad y el contexto en la recuperación de documentos.

ParentDocumentRetriever: Equilibrando Especificidad y Contexto

El ParentDocumentRetriever crea fragmentos pequeños para embeddings precisos mientras retiene suficiente contexto para una recuperación significativa. Recupera fragmentos precisos y sus documentos principales, asegurando especificidad y contexto sin perder información importante.

A continuación, veamos cómo el chunking de PDF y la indexación trabajan juntos para mejorar el rendimiento de la recuperación.

Chunking de PDF e Indexación en FAISS Vector Store

El siguiente fragmento de código demuestra cómo preparar documentos PDF para RAG dividiéndolos en chunks e indexándolos en un FAISS vector store. Este proceso involucra varios pasos, incluyendo la definición de los tamaños de los chunks, la creación de embeddings, la configuración del almacenamiento para los documentos y su indexación para búsquedas de similitud. Los tamaños de chunks principales y secundarios y los tamaños de superposición son parámetros importantes que influyen en la granularidad de los chunks y el nivel de contexto retenido, afectando así la calidad de la recuperación.

Los documentos se cargan utilizando un cargador de PDF, y se crea el ParentDocumentRetriever para manejar la recuperación de chunks de documentos junto con sus documentos principales. El código luego itera sobre los documentos cargados, agregándolos al retriever en lotes, y finalmente guarda la base de datos indexada localmente para su uso futuro.

# Constants for chunk overlap
CHILD_CHUNK_SIZE = 1024
CHILD_CHUNK_OVERLAP = 100
PARENT_CHUNK_SIZE = 4096
PARENT_CHUNK_OVERLAP = 400

# Create embeddings instance
embeddings = MistralAIEmbeddings(model="mistral-embed", mistral_api_key=my_api_key)

# Settings
index_name = args.index_name
data_files_path = args.data_files_path
dbstore_path = args.dbstore_path
docstore_path = args.docstore_path

# Create stores
fs = LocalFileStore(docstore_path)
store = create_kv_docstore(fs)

# Create Parent and Child text splitters
child_text_splitter = RecursiveCharacterTextSplitter(chunk_size=CHILD_CHUNK_SIZE, chunk_overlap=CHILD_CHUNK_OVERLAP)
parent_text_splitter = RecursiveCharacterTextSplitter(chunk_size=PARENT_CHUNK_SIZE, chunk_overlap=PARENT_CHUNK_OVERLAP)

# Create FAISS vectorstore
dimensions = len(embeddings.embed_query("dummy"))

db = FAISS(
    embedding_function=embeddings,
    index=IndexFlatL2(dimensions),
    docstore=InMemoryDocstore(),
    index_to_docstore_id={},
    normalize_L2=False
)

# Load documents
logging.info("Loading documents...")
loader = PyPDFDirectoryLoader(data_files_path)
docs = loader.load()
logging.info(f"Number of document blocks loaded: {len(docs)}")

# Create ParentDocumentRetriever
big_chunks_retriever = ParentDocumentRetriever(
    vectorstore=db,
    docstore=store,
    child_splitter=child_text_splitter,
    parent_splitter=parent_text_splitter
)

# Add documents to retriever
MAX_BATCH_SIZE = 100

for i in tqdm(range(0, len(docs), MAX_BATCH_SIZE)):
    logging.info(f"Start: {i}")
    i_end = min(len(docs), i + MAX_BATCH_SIZE)
    logging.info(f"End: {i_end}")
    batch = docs[i:i_end]
    try:
        big_chunks_retriever.add_documents(batch, ids=None)
    except ValueError as e:
        logging.error(e)
        big_chunks_retriever.add_documents(batch[:50], ids=None)
        big_chunks_retriever.add_documents(batch[50:], ids=None)
        continue
    logging.info(f"Number of keys stored in the docstore: {len(list(store.yield_keys()))}")

# Save the database
db.save_local(dbstore_path, index_name)
logging.info("Completed")

Diagrama de la Cadena Naive RAG

A continuación se muestra un diagrama que representa el flujo de trabajo del sistema Naive RAG, mostrando cómo cada componente interactúa para generar respuestas informadas.

Habiendo ilustrado el flujo de trabajo, veamos ahora cómo definir la cadena RAG utilizando LangChain Expression Language (LCEL).

Definiendo la Cadena RAG con LangChain Expression Language (LCEL)

Esta sección explica cómo configurar el Parent Document Retriever utilizando LangChain, incluyendo la configuración del vector store y la definición de los text splitters.

El primer fragmento de código demuestra cómo definir una función para reconstruir el retriever. Esta función toma rutas para el almacén de documentos y el almacén de la base de datos, así como el modelo de embeddings y el nombre del índice. Crea tanto text splitters secundarios como principales para asegurar que los documentos se dividan adecuadamente para la indexación y recuperación.

# Constants
NUM_CTX = 32768
RETRIEVED_CHUNKS = 20
CHILD_CHUNK_SIZE = 1024
CHILD_CHUNK_OVERLAP = 100
PARENT_CHUNK_SIZE = 4096
PARENT_CHUNK_OVERLAP = 400

def rebuild_retriever(docstore_path, dbstore_path, embeddings, index_name):
    child_text_splitter = RecursiveCharacterTextSplitter(chunk_size=CHILD_CHUNK_SIZE, chunk_overlap=CHILD_CHUNK_OVERLAP)
    parent_text_splitter = RecursiveCharacterTextSplitter(chunk_size=PARENT_CHUNK_SIZE, chunk_overlap=PARENT_CHUNK_OVERLAP)

    fs = LocalFileStore(docstore_path)
    docstore = create_kv_docstore(fs)

    vectordb = FAISS.load_local(
        folder_path=dbstore_path,
        embeddings=embeddings,
        index_name=index_name,
        allow_dangerous_deserialization=True
    )

    big_chunks_retriever = ParentDocumentRetriever(
        vectorstore=vectordb,
        docstore=docstore,
        child_splitter=child_text_splitter,
        parent_splitter=parent_text_splitter,
        search_type="similarity",
        search_kwargs={
            "k": RETRIEVED_CHUNKS
        }
    )
    return big_chunks_retriever

A continuación, definamos una plantilla de prompt que guiará la respuesta del modelo de lenguaje y asegurará que solo se base en el contexto recuperado.

# Define Prompt Template
def get_prompt_template() -> PromptTemplate:
    """
    Define and return a prompt template for question-answering tasks.

    Returns:
        PromptTemplate: The prompt template for question-answering tasks.
    """
    system_prompt = """
You are an assistant for question-answering tasks related to a knowledge domain based on a context provided to you.
Answer the question only based on the provided context.
If the context does not contain the information, just say that you don't know and don´t give any other response.
Give a response in technical English language and do not translate acronyms in the response.
Include the references at the end of the response, specifying only the name of the document and the page number(s) of the documents in the context used to build the response.
Do not include metadata information in the list of documents used to build the response.
Avoid duplicates in the list of documents used to build the response.

Example of output:

Response goes here

References:
- Document Name: name goes here - Page Number(s): pages go here
    """

    prompt = ChatPromptTemplate.from_messages([
        ("system", system_prompt),
        ("human", "{question}\n Context: {context}\n"),
    ])

    return prompt

Finalmente, el último fragmento muestra cómo construir toda la cadena RAG en LCEL, integrando el retriever, el prompt y el modelo de lenguaje.

# Build retriever
big_chunks_retriever = rebuild_retriever(docstore_path, dbstore_path, embeddings, index_name)

# Define prompt template
prompt = get_prompt_template()

llm = ChatMistralAI(model="mistral-large-2407", mistral_api_key=my_api_key, temperature=0.0, num_ctx=NUM_CTX)
naive_chain = (
    {
        "context": big_chunks_retriever, "question": RunnablePassthrough()
    }
    | prompt
    | llm
    | StrOutputParser()
)

Enlaces Útiles

¿Disfrutaste este post? ¿Lo encontraste útil? No dudes en dejar un comentario a continuación para compartir tus pensamientos o hacer preguntas. Se requiere una cuenta de GitHub para unirse a la discusión.