Parent Document Retriever en Action : Configuration de RAG avec Mistral LLM et LangChain

Introduction

Ce post démontre comment configurer un système de Retrieval-Augmented Generation (RAG) en utilisant LangChain, en intégrant un Parent Document Retriever avec les modèles de Mistral AI. Il fournit des détails d'implémentation, y compris du code Python, pour montrer comment RAG améliore la qualité des réponses des modèles de langage. Cela complète le post associé, qui présente le cas final.

Architecture Clé de RAG

Voici les composants clés de ce système RAG, décrivant leurs rôles et contributions à l'architecture globale :

Mistral LLM et embeddings : Utilise les embeddings de Mistral et le LLM Mistral Large pour générer des réponses pertinentes basées sur des connaissances externes.
Parent Document Retriever : Récupère de petits morceaux d'informations tout en référençant leurs documents parents pour le contexte.
FAISS Vector Store : Stocke les embeddings et permet des recherches de similarité efficaces.
Document Chunking : Divise les documents PDF en parties plus petites pour une meilleure récupération.
Naive RAG Chain : Connecte le retriever, le vector store et le LLM pour générer des réponses informées.

Avec ces composants en tête, explorons maintenant comment ils fonctionnent ensemble pour créer un système efficace de Retrieval-Augmented Generation.

Framework LangChain

LangChain est le framework principal utilisé pour implémenter le système de Retrieval-Augmented Generation (RAG). Il fournit les outils et composants nécessaires pour construire l'application RAG, y compris l'intégration avec les vector stores et les retrievers.

Maintenant que nous avons un framework, approfondissons comment le ParentDocumentRetriever équilibre spécificité et contexte dans la récupération de documents.

ParentDocumentRetriever : Équilibrer Spécificité et Contexte

Le ParentDocumentRetriever crée de petits morceaux pour des embeddings précis tout en conservant suffisamment de contexte pour une récupération significative. Il récupère des morceaux précis et leurs documents parents, assurant spécificité et contexte sans perdre d'informations importantes.

Ensuite, examinons comment le chunking de PDF et l'indexation fonctionnent ensemble pour améliorer les performances de récupération.

Chunking de PDF et Indexation dans le FAISS Vector Store

L'extrait de code suivant montre comment préparer des documents PDF pour RAG en les divisant en morceaux et en les indexant dans un FAISS vector store. Ce processus implique plusieurs étapes, y compris la définition des tailles de morceaux, la création d'embeddings, la configuration du stockage des documents et leur indexation pour des recherches de similarité. Les tailles de morceaux parents et enfants et les tailles de chevauchement sont des paramètres importants qui influencent la granularité des morceaux et le niveau de contexte conservé, affectant ainsi la qualité de la récupération.

Les documents sont chargés à l'aide d'un chargeur PDF, et le ParentDocumentRetriever est créé pour gérer la récupération des morceaux de documents ainsi que leurs documents parents. Le code parcourt ensuite les documents chargés, les ajoute au retriever par lots, et sauvegarde finalement la base de données indexée localement pour une utilisation future.

# 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")

Diagramme de la Naive RAG Chain

Ci-dessous se trouve un diagramme représentant le workflow du système Naive RAG, montrant comment chaque composant interagit pour générer des réponses informées.

Ayant illustré le workflow, examinons maintenant comment définir la RAG chain en utilisant LangChain Expression Language (LCEL).

Définition de la RAG Chain avec LangChain Expression Language (LCEL)

Cette section explique comment configurer le Parent Document Retriever en utilisant LangChain, y compris la configuration du vector store et la définition des text splitters.

Le premier extrait de code montre comment définir une fonction pour reconstruire le retriever. Cette fonction prend en entrée les chemins pour le document store et le database store, ainsi que le modèle d'embeddings et le nom de l'index. Elle crée à la fois des text splitters enfants et parents pour s'assurer que les documents sont correctement divisés pour l'indexation et la récupération.

# 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

Ensuite, définissons un modèle de prompt qui guidera la réponse du modèle de langage et s'assurera qu'il ne s'appuie que sur le contexte récupéré.

# 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

Enfin, le dernier extrait montre comment construire l'ensemble de la RAG chain en LCEL, en intégrant le retriever, le prompt et le modèle de langage.

# 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()
)

Liens Utiles

Vous avez apprécié ce post ? Vous l'avez trouvé utile ? N'hésitez pas à laisser un commentaire ci-dessous pour partager vos réflexions ou poser des questions. Un compte GitHub est requis pour participer à la discussion.