all_rag_techniques/context_enrichment_window_around_chunk.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Context Enrichment Window for Document Retrieval\n",
    "\n",
    "## Overview\n",
    "\n",
    "This code implements a context enrichment window technique for document retrieval in a vector database. It enhances the standard retrieval process by adding surrounding context to each retrieved chunk, improving the coherence and completeness of the returned information.\n",
    "\n",
    "## Motivation\n",
    "\n",
    "Traditional vector search often returns isolated chunks of text, which may lack necessary context for full understanding. This approach aims to provide a more comprehensive view of the retrieved information by including neighboring text chunks.\n",
    "\n",
    "## Key Components\n",
    "\n",
    "1. PDF processing and text chunking\n",
    "2. Vector store creation using FAISS and OpenAI embeddings\n",
    "3. Custom retrieval function with context window\n",
    "4. Comparison between standard and context-enriched retrieval\n",
    "\n",
    "## Method Details\n",
    "\n",
    "### Document Preprocessing\n",
    "\n",
    "1. The PDF is read and converted to a string.\n",
    "2. The text is split into chunks with overlap, each chunk tagged with its index.\n",
    "\n",
    "### Vector Store Creation\n",
    "\n",
    "1. OpenAI embeddings are used to create vector representations of the chunks.\n",
    "2. A FAISS vector store is created from these embeddings.\n",
    "\n",
    "### Context-Enriched Retrieval\n",
    "\n",
    "1. The `retrieve_with_context_overlap` function performs the following steps:\n",
    "   - Retrieves relevant chunks based on the query\n",
    "   - For each relevant chunk, fetches neighboring chunks\n",
    "   - Concatenates the chunks, accounting for overlap\n",
    "   - Returns the expanded context for each relevant chunk\n",
    "\n",
    "### Retrieval Comparison\n",
    "\n",
    "The notebook includes a section to compare standard retrieval with the context-enriched approach.\n",
    "\n",
    "## Benefits of this Approach\n",
    "\n",
    "1. Provides more coherent and contextually rich results\n",
    "2. Maintains the advantages of vector search while mitigating its tendency to return isolated text fragments\n",
    "3. Allows for flexible adjustment of the context window size\n",
    "\n",
    "## Conclusion\n",
    "\n",
    "This context enrichment window technique offers a promising way to improve the quality of retrieved information in vector-based document search systems. By providing surrounding context, it helps maintain the coherence and completeness of the retrieved information, potentially leading to better understanding and more accurate responses in downstream tasks such as question answering."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div style=\"text-align: center;\">\n",
    "\n",
    "<img src=\"../images/vector-search-comparison_context_enrichment.svg\" alt=\"context enrichment window\" style=\"width:70%; height:auto;\">\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div style=\"text-align: center;\">\n",
    "\n",
    "<img src=\"../images/context_enrichment_window.svg\" alt=\"context enrichment window\" style=\"width:70%; height:auto;\">\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Import libraries and environment variables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "c:\\Users\\N7\\PycharmProjects\\llm_tasks\\RAG_TECHNIQUES\\.venv\\Lib\\site-packages\\deepeval\\__init__.py:45: UserWarning: You are using deepeval version 0.21.73, however version 0.21.78 is available. You should consider upgrading via the \"pip install --upgrade deepeval\" command.\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import sys\n",
    "from dotenv import load_dotenv\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 path to PDF"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "path = \"../data/Understanding_Climate_Change.pdf\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Read PDF to string"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "content = read_pdf_to_string(path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Function to split text into chunks with metadata of the chunk chronological index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def split_text_to_chunks_with_indices(text: str, chunk_size: int, chunk_overlap: int) -> List[Document]:\n",
    "    chunks = []\n",
    "    start = 0\n",
    "    while start < len(text):\n",
    "        end = start + chunk_size\n",
    "        chunk = text[start:end]\n",
    "        chunks.append(Document(page_content=chunk, metadata={\"index\": len(chunks), \"text\": text}))\n",
    "        start += chunk_size - chunk_overlap\n",
    "    return chunks"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Split our document accordingly"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "chunks_size = 400\n",
    "chunk_overlap = 200\n",
    "docs = split_text_to_chunks_with_indices(content, chunks_size, chunk_overlap)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create vector store and retriever"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "embeddings = OpenAIEmbeddings()\n",
    "vectorstore = FAISS.from_documents(docs, embeddings)\n",
    "chunks_query_retriever = vectorstore.as_retriever(search_kwargs={\"k\": 1})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Function to draw the k<sup>th</sup> chunk (in the original order) from the vector store \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_chunk_by_index(vectorstore, target_index: int) -> Document:\n",
    "    \"\"\"\n",
    "    Retrieve a chunk from the vectorstore based on its index in the metadata.\n",
    "    \n",
    "    Args:\n",
    "    vectorstore (VectorStore): The vectorstore containing the chunks.\n",
    "    target_index (int): The index of the chunk to retrieve.\n",
    "    \n",
    "    Returns:\n",
    "    Optional[Document]: The retrieved chunk as a Document object, or None if not found.\n",
    "    \"\"\"\n",
    "    # This is a simplified version. In practice, you might need a more efficient method\n",
    "    # to retrieve chunks by index, depending on your vectorstore implementation.\n",
    "    all_docs = vectorstore.similarity_search(\"\", k=vectorstore.index.ntotal)\n",
    "    for doc in all_docs:\n",
    "        if doc.metadata.get('index') == target_index:\n",
    "            return doc\n",
    "    return None"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Check the function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Understanding Climate Change \n",
      "Chapter 1: Introduction to Climate Change \n",
      "Climate change refers to significant, long-term changes in the global climate. The term \n",
      "\"global climate\" encompasses the planet's overall weather patterns, including temperature, \n",
      "precipitation, and wind patterns, over an extended period. Over the past century, human \n",
      "activities, particularly the burning of fossil fuels and \n"
     ]
    }
   ],
   "source": [
    "chunk = get_chunk_by_index(vectorstore, 0)\n",
    "print(chunk.page_content)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Function that retrieves from the vector stroe based on semantic similarity and then pads each retrieved chunk with its num_neighbors before and after, taking into account the chunk overlap to construct a meaningful wide window arround it"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {},
   "outputs": [],
   "source": [
    "def retrieve_with_context_overlap(vectorstore, retriever, query: str, num_neighbors: int = 1, chunk_size: int = 200, chunk_overlap: int = 20) -> List[str]:\n",
    "    \"\"\"\n",
    "    Retrieve chunks based on a query, then fetch neighboring chunks and concatenate them, \n",
    "    accounting for overlap and correct indexing.\n",
    "\n",
    "    Args:\n",
    "    vectorstore (VectorStore): The vectorstore containing the chunks.\n",
    "    retriever: The retriever object to get relevant documents.\n",
    "    query (str): The query to search for relevant chunks.\n",
    "    num_neighbors (int): The number of chunks to retrieve before and after each relevant chunk.\n",
    "    chunk_size (int): The size of each chunk when originally split.\n",
    "    chunk_overlap (int): The overlap between chunks when originally split.\n",
    "\n",
    "    Returns:\n",
    "    List[str]: List of concatenated chunk sequences, each centered on a relevant chunk.\n",
    "    \"\"\"\n",
    "    relevant_chunks = retriever.get_relevant_documents(query)\n",
    "    result_sequences = []\n",
    "\n",
    "    for chunk in relevant_chunks:\n",
    "        current_index = chunk.metadata.get('index')\n",
    "        if current_index is None:\n",
    "            continue\n",
    "\n",
    "        # Determine the range of chunks to retrieve\n",
    "        start_index = max(0, current_index - num_neighbors)\n",
    "        end_index = current_index + num_neighbors + 1  # +1 because range is exclusive at the end\n",
    "\n",
    "        # Retrieve all chunks in the range\n",
    "        neighbor_chunks = []\n",
    "        for i in range(start_index, end_index):\n",
    "            neighbor_chunk = get_chunk_by_index(vectorstore, i)\n",
    "            if neighbor_chunk:\n",
    "                neighbor_chunks.append(neighbor_chunk)\n",
    "\n",
    "        # Sort chunks by their index to ensure correct order\n",
    "        neighbor_chunks.sort(key=lambda x: x.metadata.get('index', 0))\n",
    "\n",
    "        # Concatenate chunks, accounting for overlap\n",
    "        concatenated_text = neighbor_chunks[0].page_content\n",
    "        for i in range(1, len(neighbor_chunks)):\n",
    "            current_chunk = neighbor_chunks[i].page_content\n",
    "            overlap_start = max(0, len(concatenated_text) - chunk_overlap)\n",
    "            concatenated_text = concatenated_text[:overlap_start] + current_chunk\n",
    "\n",
    "        result_sequences.append(concatenated_text)\n",
    "\n",
    "    return result_sequences"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Comparing regular retrival and retrival with context window"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Baseline approach\n",
    "query = \"Explain the role of deforestation and fossil fuels in climate change.\"\n",
    "baseline_chunk = chunks_query_retriever.get_relevant_documents(query\n",
    "    ,\n",
    "    k=1\n",
    ")\n",
    "# Focused context enrichment approach\n",
    "enriched_chunks = retrieve_with_context_overlap(\n",
    "    vectorstore,\n",
    "    chunks_query_retriever,\n",
    "    query,\n",
    "    num_neighbors=1,\n",
    "    chunk_size=400,\n",
    "    chunk_overlap=200\n",
    ")\n",
    "\n",
    "print(\"Baseline Chunk:\")\n",
    "print(baseline_chunk[0].page_content)\n",
    "print(\"\\nEnriched Chunks:\")\n",
    "print(enriched_chunks[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### An example that showcases the superiority of additional context window"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Regular retrieval:\n",
      "\n",
      "Context 1:\n",
      "\n",
      "Deep Learning, a subset of machine learning using neural networks with many layers, began to show promising results in the early 2010s. The breakthrough came in 2012 when a deep neural network significantly outperformed other machine learning method\n",
      "\n",
      "\n",
      "\n",
      "Retrieval with context overlap:\n",
      "\n",
      "Context 1:\n",
      "ng multi-layer networks during this time.\n",
      "\n",
      "The late 1990s and 2000s marked the rise of machine learning approaches. Support Vector Machines (SVMs) and Random Forests became popular for various classification and regression tasks.\n",
      "\n",
      "Deep Learning, a subset of machine learning using neural networks with many layers, began to show promising results in the early 2010s. The breakthrough came in 2012 when a deep neural network significantly outperformed other machine learning methods in the ImageNet competition.\n",
      "\n",
      "Since then, deep learning has revolutionized many AI applications, including image and speech recognition, natural language processing, and game playing. In 2016, Google's AlphaGo defeated a world c\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "\n",
    "document_content = \"\"\"\n",
    "Artificial Intelligence (AI) has a rich history dating back to the mid-20th century. The term \"Artificial Intelligence\" was coined in 1956 at the Dartmouth Conference, marking the field's official beginning.\n",
    "\n",
    "In the 1950s and 1960s, AI research focused on symbolic methods and problem-solving. The Logic Theorist, created in 1955 by Allen Newell and Herbert A. Simon, is often considered the first AI program.\n",
    "\n",
    "The 1960s saw the development of expert systems, which used predefined rules to solve complex problems. DENDRAL, created in 1965, was one of the first expert systems, designed to analyze chemical compounds.\n",
    "\n",
    "However, the 1970s brought the first \"AI Winter,\" a period of reduced funding and interest in AI research, largely due to overpromised capabilities and underdelivered results.\n",
    "\n",
    "The 1980s saw a resurgence with the popularization of expert systems in corporations. The Japanese government's Fifth Generation Computer Project also spurred increased investment in AI research globally.\n",
    "\n",
    "Neural networks gained prominence in the 1980s and 1990s. The backpropagation algorithm, although discovered earlier, became widely used for training multi-layer networks during this time.\n",
    "\n",
    "The late 1990s and 2000s marked the rise of machine learning approaches. Support Vector Machines (SVMs) and Random Forests became popular for various classification and regression tasks.\n",
    "\n",
    "Deep Learning, a subset of machine learning using neural networks with many layers, began to show promising results in the early 2010s. The breakthrough came in 2012 when a deep neural network significantly outperformed other machine learning methods in the ImageNet competition.\n",
    "\n",
    "Since then, deep learning has revolutionized many AI applications, including image and speech recognition, natural language processing, and game playing. In 2016, Google's AlphaGo defeated a world champion Go player, a landmark achievement in AI.\n",
    "\n",
    "The current era of AI is characterized by the integration of deep learning with other AI techniques, the development of more efficient and powerful hardware, and the ethical considerations surrounding AI deployment.\n",
    "\n",
    "Transformers, introduced in 2017, have become a dominant architecture in natural language processing, enabling models like GPT (Generative Pre-trained Transformer) to generate human-like text.\n",
    "\n",
    "As AI continues to evolve, new challenges and opportunities arise. Explainable AI, robust and fair machine learning, and artificial general intelligence (AGI) are among the key areas of current and future research in the field.\n",
    "\"\"\"\n",
    "\n",
    "chunks_size = 250\n",
    "chunk_overlap = 20\n",
    "document_chunks = split_text_to_chunks_with_indices(document_content, chunks_size, chunk_overlap)\n",
    "document_vectorstore = FAISS.from_documents(document_chunks, embeddings)\n",
    "document_retriever = document_vectorstore.as_retriever(search_kwargs={\"k\": 1})\n",
    "\n",
    "query = \"When did deep learning become prominent in AI?\"\n",
    "context = document_retriever.get_relevant_documents(query)\n",
    "context_pages_content = [doc.page_content for doc in context]\n",
    "\n",
    "print(\"Regular retrieval:\\n\")\n",
    "show_context(context_pages_content)\n",
    "\n",
    "sequences = retrieve_with_context_overlap(document_vectorstore, document_retriever, query, num_neighbors=1)\n",
    "print(\"\\nRetrieval with context enrichment:\\n\")\n",
    "show_context(sequences)"
   ]
  }
 ],
 "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
}