run-llama--llama_index/docs/examples/node_postprocessor/AIMonRerank.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a href=\"https://colab.research.google.com/drive/1N4agIVU1NTEHaO5mLPa-bGBNpY17AxsO#scrollTo=8n2pDnW7521g\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# AIMon Rerank"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you're opening this Notebook on colab, you will probably need to install LlamaIndex 🦙"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture\n",
    "!pip install llama-index\n",
    "!pip install llama-index-postprocessor-aimon-rerank"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core import VectorStoreIndex, SimpleDirectoryReader\n",
    "from llama_index.core.response.pprint_utils import pprint_response"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "An OpenAI and AIMon API key is required for this notebook. Import the AIMon and OpenAI API keys from Colab Secrets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "# Import Colab Secrets userdata module.\n",
    "from google.colab import userdata\n",
    "\n",
    "os.environ[\"AIMON_API_KEY\"] = userdata.get(\"AIMON_API_KEY\")\n",
    "os.environ[\"OPENAI_API_KEY\"] = userdata.get(\"OPENAI_API_KEY\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Download data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--2025-03-10 18:01:07--  https://raw.githubusercontent.com/run-llama/llama_index/main/docs/examples/data/paul_graham/paul_graham_essay.txt\n",
      "Resolving raw.githubusercontent.com (raw.githubusercontent.com)... 185.199.108.133, 185.199.111.133, 185.199.110.133, ...\n",
      "Connecting to raw.githubusercontent.com (raw.githubusercontent.com)|185.199.108.133|:443... connected.\n",
      "HTTP request sent, awaiting response... 200 OK\n",
      "Length: 75042 (73K) [text/plain]\n",
      "Saving to: ‘data/paul_graham/paul_graham_essay.txt’\n",
      "\n",
      "data/paul_graham/pa 100%[===================>]  73.28K  --.-KB/s    in 0.03s   \n",
      "\n",
      "2025-03-10 18:01:08 (2.59 MB/s) - ‘data/paul_graham/paul_graham_essay.txt’ saved [75042/75042]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "!mkdir -p 'data/paul_graham/'\n",
    "!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/examples/data/paul_graham/paul_graham_essay.txt' -O 'data/paul_graham/paul_graham_essay.txt'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Generate documents and build an index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load documents\n",
    "documents = SimpleDirectoryReader(\"./data/paul_graham/\").load_data()\n",
    "\n",
    "# build index\n",
    "index = VectorStoreIndex.from_documents(documents=documents)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define a task definition for the AIMon reranker and instantiate an instance of the AIMonRerank class. The [task definition](https://docs.aimon.ai/retrieval#task-definition) serves as an explicit instruction to the system, defining what the reranking evaluation should focus on."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from llama_index.postprocessor.aimon_rerank import AIMonRerank\n",
    "\n",
    "task_definition = \"Your task is to assess the actions of an individual specified in the user query against the context documents supplied.\"\n",
    "\n",
    "aimon_rerank = AIMonRerank(\n",
    "    top_n=2,\n",
    "    api_key=userdata.get(\"AIMON_API_KEY\"),\n",
    "    task_definition=task_definition,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Directly retrieve top 2 most similar nodes (i.e., without using a reranker)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "query_engine = index.as_query_engine(similarity_top_k=2)\n",
    "response = query_engine.query(\"What did Sam Altman do in this essay?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Final Response: Sam Altman was asked to become the president of Y\n",
      "Combinator, initially declined the offer to pursue starting a startup\n",
      "focused on nuclear reactors, but eventually agreed to take over\n",
      "starting with the winter 2014 batch.\n",
      "______________________________________________________________________\n",
      "Source Node 1/2\n",
      "Node ID: 2940ea4a-69ec-4fc4-9dd4-8ed54a9d4f1b\n",
      "Similarity: 0.8305926707169754\n",
      "Text: When I was dealing with some urgent problem during YC, there was\n",
      "about a 60% chance it had to do with HN, and a 40% chance it had do\n",
      "with everything else combined. [17]  As well as HN, I wrote all of\n",
      "YC's internal software in Arc. But while I continued to work a good\n",
      "deal in Arc, I gradually stopped working on Arc, partly because I\n",
      "didn't have t...\n",
      "______________________________________________________________________\n",
      "Source Node 2/2\n",
      "Node ID: 2f043635-e4ce-4054-92f3-b624fd90ae04\n",
      "Similarity: 0.8239262662012308\n",
      "Text: Up till that point YC had been controlled by the original LLC we\n",
      "four had started. But we wanted YC to last for a long time, and to do\n",
      "that it couldn't be controlled by the founders. So if Sam said yes,\n",
      "we'd let him reorganize YC. Robert and I would retire, and Jessica and\n",
      "Trevor would become ordinary partners.  When we asked Sam if he wanted\n",
      "to...\n"
     ]
    }
   ],
   "source": [
    "pprint_response(response, show_source=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Retrieve top 10 most relevant nodes, but then rerank with AIMon Reranker"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<img src=\"https://raw.githubusercontent.com/devvratbhardwaj/images/refs/heads/main/AIMon_Reranker.svg\" alt=\"Diagram depicting working of AIMon reranker\"/>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Explanation of the reranking process:\n",
    "\n",
    "The diagram illustrates how a reranker refines document retrieval for a more accurate response.\n",
    "\n",
    "1. **Initial Retrieval (Vector DB)**:  \n",
    "   - A query is sent to the vector database.  \n",
    "   - The system retrieves the **top 10 most relevant records** based on similarity scores (`top_k = 10`).  \n",
    "\n",
    "2. **Reranking with AIMon**:  \n",
    "   - Instead of using only the highest-scoring records directly, these 10 records are reranked using the **AIMon Reranker**.  \n",
    "   - The reranker evaluates the documents based on their actual relevance to the query, rather than just raw similarity scores.  \n",
    "   - During this step, a **task definition** is applied, serving as an explicit instruction that defines what the reranking evaluation should focus on.  \n",
    "   - This ensures that the selected records are not just statistically similar but also **contextually relevant** to the intended task.  \n",
    "\n",
    "3. **Final Selection (`top_n = 2`)**:  \n",
    "   - After reranking, the system selects the **top 2 most contextually relevant records** for response generation.  \n",
    "   - The **task definition ensures** that these records align with the query’s intent, leading to a **more precise and informative response**."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total number of batches formed: 1\n",
      "Processing batch 1/1 with 10 context documents.\n",
      "Finished processing. Total batches sent to AIMon reranker: 1\n",
      "Top 2 nodes selected after reranking.\n"
     ]
    }
   ],
   "source": [
    "query_engine = index.as_query_engine(\n",
    "    similarity_top_k=10, node_postprocessors=[aimon_rerank]\n",
    ")\n",
    "response = query_engine.query(\"What did Sam Altman do in this essay?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Final Response: Sam Altman was asked to become the president of Y\n",
      "Combinator, initially declined the offer to pursue starting a startup\n",
      "focused on nuclear reactors, but eventually agreed to take over as\n",
      "president starting with the winter 2014 batch.\n",
      "______________________________________________________________________\n",
      "Source Node 1/2\n",
      "Node ID: 2940ea4a-69ec-4fc4-9dd4-8ed54a9d4f1b\n",
      "Similarity: 0.48260445005911023\n",
      "Text: When I was dealing with some urgent problem during YC, there was\n",
      "about a 60% chance it had to do with HN, and a 40% chance it had do\n",
      "with everything else combined. [17]  As well as HN, I wrote all of\n",
      "YC's internal software in Arc. But while I continued to work a good\n",
      "deal in Arc, I gradually stopped working on Arc, partly because I\n",
      "didn't have t...\n",
      "______________________________________________________________________\n",
      "Source Node 2/2\n",
      "Node ID: 0baaf5af-6e6b-4889-8407-e49d1753980c\n",
      "Similarity: 0.48151918284717965\n",
      "Text: As Jessica and I were walking home from dinner on March 11, at\n",
      "the corner of Garden and Walker streets, these three threads\n",
      "converged. Screw the VCs who were taking so long to make up their\n",
      "minds. We'd start our own investment firm and actually implement the\n",
      "ideas we'd been talking about. I'd fund it, and Jessica could quit her\n",
      "job and work for ...\n"
     ]
    }
   ],
   "source": [
    "pprint_response(response, show_source=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Conclusion\n",
    "\n",
    "The AIMon reranker, using task definition, shifted retrieval focus from general YC leadership changes to Sam Altman’s specific actions. Initially, high-similarity documents lacked his decision-making details. After reranking, lower-similarity but contextually relevant documents highlighted his reluctance and timeline, ensuring a more accurate, task-aligned response over purely similarity-based retrieval."
   ]
  }
 ],
 "metadata": {
  "colab": {
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}