all_rag_techniques/contextual_chunk_headers.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Contextual Chunk Headers (CCH)\n",
    "\n",
    "## Overview\n",
    "\n",
    "Contextual chunk headers (CCH) is a method of creating chunk headers that contain higher-level context (such as document-level or section-level context), and prepending those chunk headers to the chunks prior to embedding them. This gives the embeddings a much more accurate and complete representation of the content and meaning of the text. In our testing, this feature leads to a substantial improvement in retrieval quality. In addition to increasing the rate at which the correct information is retrieved, CCH also reduces the rate at which irrelevant results show up in the search results. This reduces the rate at which the LLM misinterprets a piece of text in downstream chat and generation applications.\n",
    "\n",
    "## Motivation\n",
    "\n",
    "Many of the problems developers face with RAG come down to this: Individual chunks oftentimes do not contain sufficient context to be properly used by the retrieval system or the LLM. This leads to the inability to answer questions and, more worryingly, hallucinations.\n",
    "\n",
    "Examples of this problem\n",
    "- Chunks oftentimes refer to their subject via implicit references and pronouns. This causes them to not be retrieved when they should be, or to not be properly understood by the LLM.\n",
    "- Individual chunks oftentimes only make sense in the context of the entire section or document, and can be misleading when read on their own.\n",
    "\n",
    "## Key Components\n",
    "\n",
    "#### Contextual chunk headers\n",
    "The idea here is to add in higher-level context to the chunk by prepending a chunk header. This chunk header could be as simple as just the document title, or it could use a combination of document title, a concise document summary, and the full hierarchy of section and sub-section titles.\n",
    "\n",
    "## Method Details\n",
    "\n",
    "#### Context generation\n",
    "In the demonstration below we use an LLM to generate a descriptive title for the document. This is done through a simple prompt where you pass in a truncated version of the document text and ask the LLM to generate a descriptive title for the document. If you already have sufficiently descriptive document titles then you can directly use those instead. We've found that a document title is the simplest and most important kind of higher-level context to include in the chunk header.\n",
    "\n",
    "Other kinds of context you can include in the chunk header:\n",
    "- Concise document summary\n",
    "- Section/sub-section title(s)\n",
    "    - This helps the retrieval system handle queries for larger sections or topics in documents.\n",
    "\n",
    "#### Embed chunks with chunk headers\n",
    "The text you embed for each chunk is simply the concatenation of the chunk header and the chunk text. If you use a reranker during retrieval, you'll want to make sure you use this same concatenation there too.\n",
    "\n",
    "#### Add chunk headers to search results\n",
    "Including the chunk headers when presenting the search results to the LLM is also beneficial as it gives the LLM more context, and makes it less likely that it misunderstands the meaning of a chunk."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![Your Technique Name](../images/contextual_chunk_headers.svg)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "You'll need a Cohere API key and an OpenAI API key for this notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cohere\n",
    "import tiktoken\n",
    "from typing import List\n",
    "from openai import OpenAI\n",
    "import os\n",
    "from dotenv import load_dotenv\n",
    "from langchain_text_splitters import RecursiveCharacterTextSplitter\n",
    "\n",
    "# Load environment variables from a .env file\n",
    "load_dotenv()\n",
    "os.environ[\"CO_API_KEY\"] = os.getenv('CO_API_KEY') # Cohere API key\n",
    "os.environ[\"OPENAI_API_KEY\"] = os.getenv('OPENAI_API_KEY') # OpenAI API key"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load the document and split it into chunks\n",
    "We'll use the basic LangChain RecursiveCharacterTextSplitter for this demo, but you can combine CCH with more sophisticated chunking methods for even better performance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.text_splitter import RecursiveCharacterTextSplitter\n",
    "\n",
    "def split_into_chunks(text: str, chunk_size: int = 800) -> list[str]:\n",
    "    \"\"\"\n",
    "    Split a given text into chunks of specified size using RecursiveCharacterTextSplitter.\n",
    "\n",
    "    Args:\n",
    "        text (str): The input text to be split into chunks.\n",
    "        chunk_size (int, optional): The maximum size of each chunk. Defaults to 800.\n",
    "\n",
    "    Returns:\n",
    "        list[str]: A list of text chunks.\n",
    "\n",
    "    Example:\n",
    "        >>> text = \"This is a sample text to be split into chunks.\"\n",
    "        >>> chunks = split_into_chunks(text, chunk_size=10)\n",
    "        >>> print(chunks)\n",
    "        ['This is a', 'sample', 'text to', 'be split', 'into', 'chunks.']\n",
    "    \"\"\"\n",
    "    text_splitter = RecursiveCharacterTextSplitter(\n",
    "        chunk_size=chunk_size,\n",
    "        chunk_overlap=0,\n",
    "        length_function=len\n",
    "    )\n",
    "    documents = text_splitter.create_documents([text])\n",
    "    return [document.page_content for document in documents]\n",
    "\n",
    "# File path for the input document\n",
    "FILE_PATH = \"../data/nike_2023_annual_report.txt\"\n",
    "\n",
    "# Read the document and split it into chunks\n",
    "with open(FILE_PATH, \"r\") as file:\n",
    "    document_text = file.read()\n",
    "\n",
    "chunks = split_into_chunks(document_text, chunk_size=800)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generate descriptive document title to use in chunk header"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "NIKE, INC. ANNUAL REPORT ON FORM 10-K\n"
     ]
    }
   ],
   "source": [
    "# Constants\n",
    "DOCUMENT_TITLE_PROMPT = \"\"\"\n",
    "INSTRUCTIONS\n",
    "What is the title of the following document?\n",
    "\n",
    "Your response MUST be the title of the document, and nothing else. DO NOT respond with anything else.\n",
    "\n",
    "{document_title_guidance}\n",
    "\n",
    "{truncation_message}\n",
    "\n",
    "DOCUMENT\n",
    "{document_text}\n",
    "\"\"\".strip()\n",
    "\n",
    "TRUNCATION_MESSAGE = \"\"\"\n",
    "Also note that the document text provided below is just the first ~{num_words} words of the document. That should be plenty for this task. Your response should still pertain to the entire document, not just the text provided below.\n",
    "\"\"\".strip()\n",
    "\n",
    "MAX_CONTENT_TOKENS = 4000\n",
    "MODEL_NAME = \"gpt-4o-mini\"\n",
    "TOKEN_ENCODER = tiktoken.encoding_for_model('gpt-3.5-turbo')\n",
    "\n",
    "def make_llm_call(chat_messages: list[dict]) -> str:\n",
    "    \"\"\"\n",
    "    Make an API call to the OpenAI language model.\n",
    "\n",
    "    Args:\n",
    "        chat_messages (list[dict]): A list of message dictionaries for the chat completion.\n",
    "\n",
    "    Returns:\n",
    "        str: The generated response from the language model.\n",
    "    \"\"\"\n",
    "    client = OpenAI(api_key=os.getenv(\"OPENAI_API_KEY\"))\n",
    "    response = client.chat.completions.create(\n",
    "        model=MODEL_NAME,\n",
    "        messages=chat_messages,\n",
    "        max_tokens=MAX_CONTENT_TOKENS,\n",
    "        temperature=0.2,\n",
    "    )\n",
    "    return response.choices[0].message.content.strip()\n",
    "\n",
    "def truncate_content(content: str, max_tokens: int) -> tuple[str, int]:\n",
    "    \"\"\"\n",
    "    Truncate the content to a specified maximum number of tokens.\n",
    "\n",
    "    Args:\n",
    "        content (str): The input text to be truncated.\n",
    "        max_tokens (int): The maximum number of tokens to keep.\n",
    "\n",
    "    Returns:\n",
    "        tuple[str, int]: A tuple containing the truncated content and the number of tokens.\n",
    "    \"\"\"\n",
    "    tokens = TOKEN_ENCODER.encode(content, disallowed_special=())\n",
    "    truncated_tokens = tokens[:max_tokens]\n",
    "    return TOKEN_ENCODER.decode(truncated_tokens), min(len(tokens), max_tokens)\n",
    "\n",
    "def get_document_title(document_text: str, document_title_guidance: str = \"\") -> str:\n",
    "    \"\"\"\n",
    "    Extract the title of a document using a language model.\n",
    "\n",
    "    Args:\n",
    "        document_text (str): The text of the document.\n",
    "        document_title_guidance (str, optional): Additional guidance for title extraction. Defaults to \"\".\n",
    "\n",
    "    Returns:\n",
    "        str: The extracted document title.\n",
    "    \"\"\"\n",
    "    # Truncate the content if it's too long\n",
    "    document_text, num_tokens = truncate_content(document_text, MAX_CONTENT_TOKENS)\n",
    "    truncation_message = TRUNCATION_MESSAGE.format(num_words=3000) if num_tokens >= MAX_CONTENT_TOKENS else \"\"\n",
    "\n",
    "    # Prepare the prompt for title extraction\n",
    "    prompt = DOCUMENT_TITLE_PROMPT.format(\n",
    "        document_title_guidance=document_title_guidance,\n",
    "        document_text=document_text,\n",
    "        truncation_message=truncation_message\n",
    "    )\n",
    "    chat_messages = [{\"role\": \"user\", \"content\": prompt}]\n",
    "    \n",
    "    return make_llm_call(chat_messages)\n",
    "\n",
    "# Example usage\n",
    "if __name__ == \"__main__\":\n",
    "    # Assuming document_text is defined elsewhere\n",
    "    document_title = get_document_title(document_text)\n",
    "    print(f\"Document Title: {document_title}\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Add chunk header and measure impact\n",
    "Let's look at a specific example to demonstrate the impact of adding a chunk header. We'll use the Cohere reranker to measure relevance to a query with and without a chunk header."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Chunk header:\n",
      "Document Title: NIKE, INC. ANNUAL REPORT ON FORM 10-K\n",
      "\n",
      "Chunk text:\n",
      "Given the broad and global scope of our operations, we are particularly vulnerable to the physical risks of climate change, such \n",
      "as shifts in weather patterns. Extreme weather conditions in the areas in which our retail stores, suppliers, manufacturers, \n",
      "customers, distribution centers, offices, headquarters and vendors are located could adversely affect our operating results and \n",
      "financial condition. Moreover, natural disasters such as earthquakes, hurricanes, wildfires, tsunamis, floods or droughts, whether \n",
      "occurring in the United States or abroad, and their related consequences and effects, including energy shortages and public \n",
      "health issues, have in the past temporarily disrupted, and could in the future disrupt, our operations, the operations of our\n",
      "\n",
      "Query: Nike climate change impact\n",
      "\n",
      "Similarity w/o contextual chunk header: 0.10576342\n",
      "Similarity with contextual chunk header: 0.92206234\n"
     ]
    }
   ],
   "source": [
    "def rerank_documents(query: str, chunks: List[str]) -> List[float]:\n",
    "    \"\"\"\n",
    "    Use Cohere Rerank API to rerank the search results.\n",
    "\n",
    "    Args:\n",
    "        query (str): The search query.\n",
    "        chunks (List[str]): List of document chunks to be reranked.\n",
    "\n",
    "    Returns:\n",
    "        List[float]: List of similarity scores for each chunk, in the original order.\n",
    "    \"\"\"\n",
    "    MODEL = \"rerank-english-v3.0\"\n",
    "    client = cohere.Client(api_key=os.environ[\"CO_API_KEY\"])\n",
    "\n",
    "    reranked_results = client.rerank(model=MODEL, query=query, documents=chunks)\n",
    "    results = reranked_results.results\n",
    "    reranked_indices = [result.index for result in results]\n",
    "    reranked_similarity_scores = [result.relevance_score for result in results]\n",
    "    \n",
    "    # Convert back to order of original documents\n",
    "    similarity_scores = [0] * len(chunks)\n",
    "    for i, index in enumerate(reranked_indices):\n",
    "        similarity_scores[index] = reranked_similarity_scores[i]\n",
    "\n",
    "    return similarity_scores\n",
    "\n",
    "def compare_chunk_similarities(chunk_index: int, chunks: List[str], document_title: str, query: str) -> None:\n",
    "    \"\"\"\n",
    "    Compare similarity scores for a chunk with and without a contextual header.\n",
    "\n",
    "    Args:\n",
    "        chunk_index (int): Index of the chunk to inspect.\n",
    "        chunks (List[str]): List of all document chunks.\n",
    "        document_title (str): Title of the document.\n",
    "        query (str): The search query to use for comparison.\n",
    "\n",
    "    Prints:\n",
    "        Chunk header, chunk text, query, and similarity scores with and without the header.\n",
    "    \"\"\"\n",
    "    chunk_text = chunks[chunk_index]\n",
    "    chunk_wo_header = chunk_text\n",
    "    chunk_w_header = f\"Document Title: {document_title}\\n\\n{chunk_text}\"\n",
    "\n",
    "    similarity_scores = rerank_documents(query, [chunk_wo_header, chunk_w_header])\n",
    "\n",
    "    print(f\"\\nChunk header:\\nDocument Title: {document_title}\")\n",
    "    print(f\"\\nChunk text:\\n{chunk_text}\")\n",
    "    print(f\"\\nQuery: {query}\")\n",
    "    print(f\"\\nSimilarity without contextual chunk header: {similarity_scores[0]:.4f}\")\n",
    "    print(f\"Similarity with contextual chunk header: {similarity_scores[1]:.4f}\")\n",
    "\n",
    "# Notebook cell for execution\n",
    "# Assuming chunks and document_title are defined in previous cells\n",
    "CHUNK_INDEX_TO_INSPECT = 86\n",
    "QUERY = \"Nike climate change impact\"\n",
    "\n",
    "compare_chunk_similarities(CHUNK_INDEX_TO_INSPECT, chunks, document_title, QUERY)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This chunk is clearly about the impact of climate change on some organization, but it doesn't explicitly say \"Nike\" in it. So the relevance to the query \"Nike climate change impact\" in only about 0.1. By simply adding the document title to the chunk that similarity goes up to 0.92."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Eval results\n",
    "\n",
    "#### KITE\n",
    "\n",
    "We evaluated CCH on an end-to-end RAG benchmark we created, called KITE (Knowledge-Intensive Task Evaluation).\n",
    "\n",
    "KITE currently consists of 4 datasets and a total of 50 questions.\n",
    "- **AI Papers** - ~100 academic papers about AI and RAG, downloaded from arXiv in PDF form.\n",
    "- **BVP Cloud 10-Ks** - 10-Ks for all companies in the Bessemer Cloud Index (~70 of them), in PDF form.\n",
    "- **Sourcegraph Company Handbook** - ~800 markdown files, with their original directory structure, downloaded from Sourcegraph's publicly accessible company handbook GitHub [page](https://github.com/sourcegraph/handbook/tree/main/content).\n",
    "- **Supreme Court Opinions** - All Supreme Court opinions from Term Year 2022 (delivered from January '23 to June '23), downloaded from the official Supreme Court [website](https://www.supremecourt.gov/opinions/slipopinion/22) in PDF form.\n",
    "\n",
    "Ground truth answers are included with each sample. Most samples also include grading rubrics. Grading is done on a scale of 0-10 for each question, with a strong LLM doing the grading.\n",
    "\n",
    "We compare performance with and without CCH. For the CCH config we use document title and document summary. All other parameters remain the same between the two configurations. We use the Cohere 3 reranker, and we use GPT-4o for response generation.\n",
    "\n",
    "|                         | No-CCH   | CCH          |\n",
    "|-------------------------|----------|--------------|\n",
    "| AI Papers               | 4.5      | 4.7          |\n",
    "| BVP Cloud               | 2.6      | 6.3          |\n",
    "| Sourcegraph             | 5.7      | 5.8          |\n",
    "| Supreme Court Opinions  | 6.1      | 7.4          |\n",
    "| **Average**             | 4.72     | 6.04         |\n",
    "\n",
    "We can see that CCH leads to an improvement in performance on each of the four datasets. Some datasets see a large improvement while others see a small improvement. The overall average score increases from 4.72 -> 6.04, a 27.9% increase.\n",
    "\n",
    "#### FinanceBench\n",
    "\n",
    "We've also evaluated CCH on FinanceBench, where it contributed to a score of 83%, compared to a baseline score of 19%. For that benchmark, we tested CCH and relevant segment extraction (RSE) jointly, so we can't say exactly how much CCH contributed to that result. But the combination of CCH and RSE clearly leads to substantial accuracy improvements on FinanceBench."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "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.9.12"
  },
  "vscode": {
   "interpreter": {
    "hash": "44d0561a9d33f22b2e67e0485c48036e39d1c698628b030a9859974b559ff507"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}