{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# Hierarchical Indices in Document Retrieval\n", "\n", "## Overview\n", "\n", "This code implements a Hierarchical Indexing system for document retrieval, utilizing two levels of encoding: document-level summaries and detailed chunks. This approach aims to improve the efficiency and relevance of information retrieval by first identifying relevant document sections through summaries, then drilling down to specific details within those sections.\n", "\n", "## Motivation\n", "\n", "Traditional flat indexing methods can struggle with large documents or corpus, potentially missing context or returning irrelevant information. Hierarchical indexing addresses this by creating a two-tier search system, allowing for more efficient and context-aware retrieval.\n", "\n", "## Key Components\n", "\n", "1. PDF processing and text chunking\n", "2. Asynchronous document summarization using OpenAI's GPT-4\n", "3. Vector store creation for both summaries and detailed chunks using FAISS and OpenAI embeddings\n", "4. Custom hierarchical retrieval function\n", "\n", "## Method Details\n", "\n", "### Document Preprocessing and Encoding\n", "\n", "1. The PDF is loaded and split into documents (likely by page).\n", "2. Each document is summarized asynchronously using GPT-4.\n", "3. The original documents are also split into smaller, detailed chunks.\n", "4. Two separate vector stores are created:\n", " - One for document-level summaries\n", " - One for detailed chunks\n", "\n", "### Asynchronous Processing and Rate Limiting\n", "\n", "1. The code uses asynchronous programming (asyncio) to improve efficiency.\n", "2. Implements batching and exponential backoff to handle API rate limits.\n", "\n", "### Hierarchical Retrieval\n", "\n", "The `retrieve_hierarchical` function implements the two-tier search:\n", "\n", "1. It first searches the summary vector store to identify relevant document sections.\n", "2. For each relevant summary, it then searches the detailed chunk vector store, filtering by the corresponding page number.\n", "3. This approach ensures that detailed information is retrieved only from the most relevant document sections.\n", "\n", "## Benefits of this Approach\n", "\n", "1. Improved Retrieval Efficiency: By first searching summaries, the system can quickly identify relevant document sections without processing all detailed chunks.\n", "2. Better Context Preservation: The hierarchical approach helps maintain the broader context of retrieved information.\n", "3. Scalability: This method is particularly beneficial for large documents or corpus, where flat searching might be inefficient or miss important context.\n", "4. Flexibility: The system allows for adjusting the number of summaries and chunks retrieved, enabling fine-tuning for different use cases.\n", "\n", "## Implementation Details\n", "\n", "1. Asynchronous Programming: Utilizes Python's asyncio for efficient I/O operations and API calls.\n", "2. Rate Limit Handling: Implements batching and exponential backoff to manage API rate limits effectively.\n", "3. Persistent Storage: Saves the generated vector stores locally to avoid unnecessary recomputation.\n", "\n", "## Conclusion\n", "\n", "Hierarchical indexing represents a sophisticated approach to document retrieval, particularly suitable for large or complex document sets. By leveraging both high-level summaries and detailed chunks, it offers a balance between broad context understanding and specific information retrieval. This method has potential applications in various fields requiring efficient and context-aware information retrieval, such as legal document analysis, academic research, or large-scale content management systems." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "
\n", "\n", "\"hierarchical_indices\"\n", "
" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Import libraries " ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import asyncio\n", "import os\n", "import sys\n", "from dotenv import load_dotenv\n", "from langchain_openai import ChatOpenAI\n", "from langchain.chains.summarize.chain import load_summarize_chain\n", "from langchain.docstore.document import Document\n", "\n", "\n", "sys.path.append(os.path.abspath(os.path.join(os.getcwd(), '..'))) # Add the parent directory to the path sicnce we work with notebooks\n", "from helper_functions import *\n", "from evaluation.evalute_rag import *\n", "\n", "# Load environment variables from a .env file\n", "load_dotenv()\n", "\n", "# Set the OpenAI API key environment variable\n", "os.environ[\"OPENAI_API_KEY\"] = os.getenv('OPENAI_API_KEY')" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Define document path" ] }, { "cell_type": "code", "execution_count": 2, "metadata": {}, "outputs": [], "source": [ "path = \"../data/Understanding_Climate_Change.pdf\"" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Function to encode to both summary and chunk levels, sharing the page metadata" ] }, { "cell_type": "code", "execution_count": 3, "metadata": {}, "outputs": [], "source": [ "async def encode_pdf_hierarchical(path, chunk_size=1000, chunk_overlap=200):\n", " \"\"\"\n", " Asynchronously encodes a PDF book into a hierarchical vector store using OpenAI embeddings.\n", " Includes rate limit handling with exponential backoff.\n", " \n", " Args:\n", " path: The path to the PDF file.\n", " chunk_size: The desired size of each text chunk.\n", " chunk_overlap: The amount of overlap between consecutive chunks.\n", " \n", " Returns:\n", " A tuple containing two FAISS vector stores:\n", " 1. Document-level summaries\n", " 2. Detailed chunks\n", " \"\"\"\n", " \n", " # Load PDF documents\n", " loader = PyPDFLoader(path)\n", " documents = await asyncio.to_thread(loader.load)\n", "\n", " # Create document-level summaries\n", " summary_llm = ChatOpenAI(temperature=0, model_name=\"gpt-4o-mini\", max_tokens=4000)\n", " summary_chain = load_summarize_chain(summary_llm, chain_type=\"map_reduce\")\n", " \n", " async def summarize_doc(doc):\n", " \"\"\"\n", " Summarizes a single document with rate limit handling.\n", " \n", " Args:\n", " doc: The document to be summarized.\n", " \n", " Returns:\n", " A summarized Document object.\n", " \"\"\"\n", " # Retry the summarization with exponential backoff\n", " summary_output = await retry_with_exponential_backoff(summary_chain.ainvoke([doc]))\n", " summary = summary_output['output_text']\n", " return Document(\n", " page_content=summary,\n", " metadata={\"source\": path, \"page\": doc.metadata[\"page\"], \"summary\": True}\n", " )\n", "\n", " # Process documents in smaller batches to avoid rate limits\n", " batch_size = 5 # Adjust this based on your rate limits\n", " summaries = []\n", " for i in range(0, len(documents), batch_size):\n", " batch = documents[i:i+batch_size]\n", " batch_summaries = await asyncio.gather(*[summarize_doc(doc) for doc in batch])\n", " summaries.extend(batch_summaries)\n", " await asyncio.sleep(1) # Short pause between batches\n", "\n", " # Split documents into detailed chunks\n", " text_splitter = RecursiveCharacterTextSplitter(\n", " chunk_size=chunk_size, chunk_overlap=chunk_overlap, length_function=len\n", " )\n", " detailed_chunks = await asyncio.to_thread(text_splitter.split_documents, documents)\n", "\n", " # Update metadata for detailed chunks\n", " for i, chunk in enumerate(detailed_chunks):\n", " chunk.metadata.update({\n", " \"chunk_id\": i,\n", " \"summary\": False,\n", " \"page\": int(chunk.metadata.get(\"page\", 0))\n", " })\n", "\n", " # Create embeddings\n", " embeddings = OpenAIEmbeddings()\n", "\n", " # Create vector stores asynchronously with rate limit handling\n", " async def create_vectorstore(docs):\n", " \"\"\"\n", " Creates a vector store from a list of documents with rate limit handling.\n", " \n", " Args:\n", " docs: The list of documents to be embedded.\n", " \n", " Returns:\n", " A FAISS vector store containing the embedded documents.\n", " \"\"\"\n", " return await retry_with_exponential_backoff(\n", " asyncio.to_thread(FAISS.from_documents, docs, embeddings)\n", " )\n", "\n", " # Generate vector stores for summaries and detailed chunks concurrently\n", " summary_vectorstore, detailed_vectorstore = await asyncio.gather(\n", " create_vectorstore(summaries),\n", " create_vectorstore(detailed_chunks)\n", " )\n", "\n", " return summary_vectorstore, detailed_vectorstore" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Encode the PDF book to both document-level summaries and detailed chunks if the vector stores do not exist\n" ] }, { "cell_type": "code", "execution_count": 4, "metadata": {}, "outputs": [], "source": [ "if os.path.exists(\"../vector_stores/summary_store\") and os.path.exists(\"../vector_stores/detailed_store\"):\n", " embeddings = OpenAIEmbeddings()\n", " summary_store = FAISS.load_local(\"../vector_stores/summary_store\", embeddings, allow_dangerous_deserialization=True)\n", " detailed_store = FAISS.load_local(\"../vector_stores/detailed_store\", embeddings, allow_dangerous_deserialization=True)\n", "\n", "else:\n", " summary_store, detailed_store = await encode_pdf_hierarchical(path)\n", " summary_store.save_local(\"../vector_stores/summary_store\")\n", " detailed_store.save_local(\"../vector_stores/detailed_store\")\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Retrieve information according to summary level, and then retrieve information from the chunk level vector store and filter according to the summary level pages" ] }, { "cell_type": "code", "execution_count": 5, "metadata": {}, "outputs": [], "source": [ "def retrieve_hierarchical(query, summary_vectorstore, detailed_vectorstore, k_summaries=3, k_chunks=5):\n", " \"\"\"\n", " Performs a hierarchical retrieval using the query.\n", "\n", " Args:\n", " query: The search query.\n", " summary_vectorstore: The vector store containing document summaries.\n", " detailed_vectorstore: The vector store containing detailed chunks.\n", " k_summaries: The number of top summaries to retrieve.\n", " k_chunks: The number of detailed chunks to retrieve per summary.\n", "\n", " Returns:\n", " A list of relevant detailed chunks.\n", " \"\"\"\n", " \n", " # Retrieve top summaries\n", " top_summaries = summary_vectorstore.similarity_search(query, k=k_summaries)\n", " \n", " relevant_chunks = []\n", " for summary in top_summaries:\n", " # For each summary, retrieve relevant detailed chunks\n", " page_number = summary.metadata[\"page\"]\n", " page_filter = lambda metadata: metadata[\"page\"] == page_number\n", " page_chunks = detailed_vectorstore.similarity_search(\n", " query, \n", " k=k_chunks, \n", " filter=page_filter\n", " )\n", " relevant_chunks.extend(page_chunks)\n", " \n", " return relevant_chunks" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Demonstrate on a use case" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "query = \"What is the greenhouse effect?\"\n", "results = retrieve_hierarchical(query, summary_store, detailed_store)\n", "\n", "# Print results\n", "for chunk in results:\n", " print(f\"Page: {chunk.metadata['page']}\")\n", " print(f\"Content: {chunk.page_content}...\") # Print first 100 characters\n", " print(\"---\")" ] } ], "metadata": { "kernelspec": { "display_name": ".venv", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.12.0" } }, "nbformat": 4, "nbformat_minor": 2 }