all_rag_techniques/reranking_with_llamaindex.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Reranking Methods in RAG Systems\n",
    "\n",
    "## Overview\n",
    "Reranking is a crucial step in Retrieval-Augmented Generation (RAG) systems that aims to improve the relevance and quality of retrieved documents. It involves reassessing and reordering initially retrieved documents to ensure that the most pertinent information is prioritized for subsequent processing or presentation.\n",
    "\n",
    "## Motivation\n",
    "The primary motivation for reranking in RAG systems is to overcome limitations of initial retrieval methods, which often rely on simpler similarity metrics. Reranking allows for more sophisticated relevance assessment, taking into account nuanced relationships between queries and documents that might be missed by traditional retrieval techniques. This process aims to enhance the overall performance of RAG systems by ensuring that the most relevant information is used in the generation phase.\n",
    "\n",
    "## Key Components\n",
    "Reranking systems typically include the following components:\n",
    "\n",
    "1. Initial Retriever: Often a vector store using embedding-based similarity search.\n",
    "2. Reranking Model: This can be either:\n",
    "   - A Large Language Model (LLM) for scoring relevance\n",
    "   - A Cross-Encoder model specifically trained for relevance assessment\n",
    "3. Scoring Mechanism: A method to assign relevance scores to documents\n",
    "4. Sorting and Selection Logic: To reorder documents based on new scores\n",
    "\n",
    "## Method Details\n",
    "The reranking process generally follows these steps:\n",
    "\n",
    "1. Initial Retrieval: Fetch an initial set of potentially relevant documents.\n",
    "2. Pair Creation: Form query-document pairs for each retrieved document.\n",
    "3. Scoring: \n",
    "   - LLM Method: Use prompts to ask the LLM to rate document relevance.\n",
    "   - Cross-Encoder Method: Feed query-document pairs directly into the model.\n",
    "4. Score Interpretation: Parse and normalize the relevance scores.\n",
    "5. Reordering: Sort documents based on their new relevance scores.\n",
    "6. Selection: Choose the top K documents from the reordered list.\n",
    "\n",
    "## Benefits of this Approach\n",
    "Reranking offers several advantages:\n",
    "\n",
    "1. Improved Relevance: By using more sophisticated models, reranking can capture subtle relevance factors.\n",
    "2. Flexibility: Different reranking methods can be applied based on specific needs and resources.\n",
    "3. Enhanced Context Quality: Providing more relevant documents to the RAG system improves the quality of generated responses.\n",
    "4. Reduced Noise: Reranking helps filter out less relevant information, focusing on the most pertinent content.\n",
    "\n",
    "## Conclusion\n",
    "Reranking is a powerful technique in RAG systems that significantly enhances the quality of retrieved information. Whether using LLM-based scoring or specialized Cross-Encoder models, reranking allows for more nuanced and accurate assessment of document relevance. This improved relevance translates directly to better performance in downstream tasks, making reranking an essential component in advanced RAG implementations.\n",
    "\n",
    "The choice between LLM-based and Cross-Encoder reranking methods depends on factors such as required accuracy, available computational resources, and specific application needs. Both approaches offer substantial improvements over basic retrieval methods and contribute to the overall effectiveness of RAG systems."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div style=\"text-align: center;\">\n",
    "\n",
    "<img src=\"../images/reranking-visualization.svg\" alt=\"rerank llm\" style=\"width:100%; height:auto;\">\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div style=\"text-align: center;\">\n",
    "\n",
    "<img src=\"../images/reranking_comparison.svg\" alt=\"rerank llm\" style=\"width:100%; height:auto;\">\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Import relevant libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "from dotenv import load_dotenv\n",
    "from typing import List\n",
    "from llama_index.core import Document\n",
    "from llama_index.core import Settings\n",
    "from llama_index.embeddings.openai import OpenAIEmbedding\n",
    "from llama_index.llms.openai import OpenAI\n",
    "from llama_index.core.readers import SimpleDirectoryReader\n",
    "from llama_index.vector_stores.faiss import FaissVectorStore\n",
    "from llama_index.core.ingestion import IngestionPipeline\n",
    "from llama_index.core.node_parser import SentenceSplitter\n",
    "from llama_index.core import VectorStoreIndex\n",
    "from llama_index.core.postprocessor import SentenceTransformerRerank, LLMRerank\n",
    "from llama_index.core import QueryBundle\n",
    "import faiss\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",
    "\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')\n",
    "\n",
    "# Llamaindex global settings for llm and embeddings\n",
    "EMBED_DIMENSION=512\n",
    "Settings.llm = OpenAI(model=\"gpt-3.5-turbo\")\n",
    "Settings.embed_model = OpenAIEmbedding(model=\"text-embedding-3-small\", dimensions=EMBED_DIMENSION)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Read docs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "path = \"../data/\"\n",
    "reader = SimpleDirectoryReader(input_dir=path, required_exts=['.pdf'])\n",
    "documents = reader.load_data()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create a vector store"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create FaisVectorStore to store embeddings\n",
    "fais_index = faiss.IndexFlatL2(EMBED_DIMENSION)\n",
    "vector_store = FaissVectorStore(faiss_index=fais_index)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Ingestion Pipeline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "base_pipeline = IngestionPipeline(\n",
    "    transformations=[SentenceSplitter()],\n",
    "    vector_store=vector_store,\n",
    "    documents=documents\n",
    ")\n",
    "\n",
    "nodes = base_pipeline.run()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Querying"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Method 1: LLM based reranking the retrieved documents"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div style=\"text-align: center;\">\n",
    "\n",
    "<img src=\"../images/rerank_llm.svg\" alt=\"rerank llm\" style=\"width:40%; height:auto;\">\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create vector index from base nodes\n",
    "index = VectorStoreIndex(nodes)\n",
    "\n",
    "query_engine_w_llm_rerank = index.as_query_engine(\n",
    "    similarity_top_k=10,\n",
    "    node_postprocessors=[\n",
    "        LLMRerank(\n",
    "            top_n=5\n",
    "        )\n",
    "    ],\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "resp = query_engine_w_llm_rerank.query(\"What are the impacts of climate change on biodiversity?\")\n",
    "print(resp)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Example that demonstrates why we should use reranking "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Comparison of Retrieval Techniques\n",
      "==================================\n",
      "Query: what is the capital of france?\n",
      "\n",
      "Baseline Retrieval Result:\n",
      "\n",
      "Document 1:\n",
      "The capital of France is great.\n",
      "\n",
      "Document 2:\n",
      "The capital of France is huge.\n",
      "\n",
      "Advanced Retrieval Result:\n",
      "\n",
      "Document 1:\n",
      "Have you ever visited Paris? It is a beautiful city where you can eat delicious food and see the Eiffel Tower. I really enjoyed all the cities in france, but its capital with the Eiffel Tower is my favorite city.\n",
      "\n",
      "Document 2:\n",
      "I really enjoyed my trip to Paris, France. The city is beautiful and the food is delicious. I would love to visit again. Such a great capital city.\n"
     ]
    }
   ],
   "source": [
    "chunks = [\n",
    "    \"The capital of France is great.\",\n",
    "    \"The capital of France is huge.\",\n",
    "    \"The capital of France is beautiful.\",\n",
    "    \"\"\"Have you ever visited Paris? It is a beautiful city where you can eat delicious food and see the Eiffel Tower. I really enjoyed all the cities in france, but its capital with the Eiffel Tower is my favorite city.\"\"\", \n",
    "    \"I really enjoyed my trip to Paris, France. The city is beautiful and the food is delicious. I would love to visit again. Such a great capital city.\"\n",
    "]\n",
    "docs = [Document(page_content=sentence) for sentence in chunks]\n",
    "\n",
    "\n",
    "def compare_rag_techniques(query: str, docs: List[Document] = docs) -> None:\n",
    "    docs = [Document(text=sentence) for sentence in chunks]\n",
    "    index = VectorStoreIndex.from_documents(docs)\n",
    "    \n",
    "    \n",
    "    print(\"Comparison of Retrieval Techniques\")\n",
    "    print(\"==================================\")\n",
    "    print(f\"Query: {query}\\n\")\n",
    "    \n",
    "    print(\"Baseline Retrieval Result:\")\n",
    "    baseline_docs = index.as_retriever(similarity_top_k=5).retrieve(query)\n",
    "    for i, doc in enumerate(baseline_docs[:2]): # Get only the first two retrieved docs\n",
    "        print(f\"\\nDocument {i+1}:\")\n",
    "        print(doc.text)\n",
    "\n",
    "    print(\"\\nAdvanced Retrieval Result:\")\n",
    "    reranker = LLMRerank(\n",
    "        top_n=2,\n",
    "    )\n",
    "    advanced_docs = reranker.postprocess_nodes(\n",
    "            baseline_docs, \n",
    "            QueryBundle(query)\n",
    "        )\n",
    "    for i, doc in enumerate(advanced_docs):\n",
    "        print(f\"\\nDocument {i+1}:\")\n",
    "        print(doc.text)\n",
    "\n",
    "\n",
    "query = \"what is the capital of france?\"\n",
    "compare_rag_techniques(query, docs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Method 2: Cross Encoder models"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div style=\"text-align: center;\">\n",
    "\n",
    "<img src=\"../images/rerank_cross_encoder.svg\" alt=\"rerank cross encoder\" style=\"width:40%; height:auto;\">\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "LlamaIndex has builtin support for [SBERT](https://www.sbert.net/index.html) models that can be used directly as node postprocessor."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "query_engine_w_cross_encoder = index.as_query_engine(\n",
    "    similarity_top_k=10,\n",
    "    node_postprocessors=[\n",
    "        SentenceTransformerRerank(\n",
    "            model='cross-encoder/ms-marco-MiniLM-L-6-v2',\n",
    "            top_n=5\n",
    "        )\n",
    "    ],\n",
    ")\n",
    "\n",
    "resp = query_engine_w_cross_encoder.query(\"What are the impacts of climate change on biodiversity?\")\n",
    "print(resp)"
   ]
  }
 ],
 "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.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}