reciprocal_rerank_fusion_algorithm.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kCkLtMsQKNXd"
      },
      "source": [
        "# Reciprocal Rerank Fusion Retriever\n",
        "\n",
        "Combining retrieval results from multiple queries and multiple indexes.\n",
        "\n",
        "The retrieved nodes will be reranked according to the Reciprocal Rerank Fusion algorithm. It provides an effecient method for rerranking retrieval results without excessive computation or reliance on external models."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3g_jWUMwKNXe"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "import openai\n",
        "\n",
        "os.environ[\"OPENAI_API_KEY\"] = \"sk-...\"\n",
        "openai.api_key = os.environ[\"OPENAI_API_KEY\"]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2BUfMe9LKNXg"
      },
      "source": [
        "## Setup"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yqqNoas4KNXh"
      },
      "source": [
        "Download Data"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Hg2Y6OWmKNXh"
      },
      "outputs": [],
      "source": [
        "!mkdir -p 'data/essay/'\n",
        "!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/examples/data/paul_graham/paul_graham_essay.txt' -O 'data/essay/paul_graham_essay.txt'"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cf_4R5X5KNXi"
      },
      "outputs": [],
      "source": [
        "from llama_index import SimpleDirectoryReader\n",
        "\n",
        "documents = SimpleDirectoryReader(\"./data/essay/\").load_data()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "u8Q2FBRBKNXi"
      },
      "source": [
        "Next, we will setup a vector index over the documentation."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tB_psH1fKNXj"
      },
      "outputs": [],
      "source": [
        "from llama_index import VectorStoreIndex, ServiceContext\n",
        "\n",
        "service_context = ServiceContext.from_defaults(chunk_size=256)\n",
        "\n",
        "index = VectorStoreIndex.from_documents(\n",
        "    documents, service_context=service_context\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UNGQRbRHKNXj"
      },
      "source": [
        "## Create a Hybrid Fusion Retriever\n",
        "\n",
        "Enables us to capture both semantic relations and keywords in our input queries.\n",
        "\n",
        "Since both of these retrievers calculate a score, we can use the reciprocal rerank algorithm to re-sort our nodes without using an additional models or excessive computation.\n",
        "\n",
        "This setup will also query 4 times, once with your original query, and generate 3 more queries.\n",
        "\n",
        "```python\n",
        "QUERY_GEN_PROMPT = (\n",
        "    \"You are a helpful assistant that generates multiple search queries based on a \"\n",
        "    \"single input query. Generate {num_queries} search queries, one on each line, \"\n",
        "    \"related to the following input query:\\n\"\n",
        "    \"Query: {query}\\n\"\n",
        "    \"Queries:\\n\"\n",
        ")\n",
        "```"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fXFDrSdlKNXj"
      },
      "source": [
        "First, we create our retrievers. Each will retrieve the top-2 most similar nodes:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1FYEklPfKNXj"
      },
      "outputs": [],
      "source": [
        "from llama_index.retrievers import BM25Retriever\n",
        "\n",
        "vector_retriever = index.as_retriever(similarity_top_k=2)\n",
        "\n",
        "bm25_retriever = BM25Retriever.from_defaults(\n",
        "    docstore=index.docstore, similarity_top_k=2\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0OzMqvW8KNXk"
      },
      "source": [
        "Creating the fusion retriever, which well return the top-2 most similar nodes from the 4 returned nodes from the retrievers:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IDYbG7_8KNXk"
      },
      "outputs": [],
      "source": [
        "from llama_index.retrievers import QueryFusionRetriever\n",
        "\n",
        "retriever = QueryFusionRetriever(\n",
        "    [vector_retriever, bm25_retriever],\n",
        "    similarity_top_k=2,\n",
        "    num_queries=4,  # set this to 1 to disable query generation\n",
        "    mode=\"reciprocal_rerank\",\n",
        "    use_async=True,\n",
        "    verbose=True,\n",
        "    # query_gen_prompt=\"...\",  # we could override the query generation prompt here\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Z2hy-9OdKNXk"
      },
      "outputs": [],
      "source": [
        "# apply nested async to run in a notebook\n",
        "import nest_asyncio\n",
        "\n",
        "nest_asyncio.apply()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tJeKpDKiKNXk",
        "outputId": "87f1880a-c0b0-4d26-8923-e0c2970d6274"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Generated queries:\n",
            "1. What were the major events or milestones in the history of Interleafe and Viaweb?\n",
            "2. Can you provide a timeline of the key developments and achievements of Interleafe and Viaweb?\n",
            "3. What were the successes and failures of Interleafe and Viaweb as companies?\n"
          ]
        }
      ],
      "source": [
        "nodes_with_scores = retriever.retrieve(\n",
        "    \"What happened at Interleafe and Viaweb?\"\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "anYDTxuhKNXl",
        "outputId": "ee047f86-7366-4a92-ce4b-7bdb62803aca"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Score: 0.05 - Now you could just update the software right on the server.\n",
            "\n",
            "We started a new company we called Viaweb, after the fact that our software worked via the web, and we got $10,000 in seed funding from Idelle's husband Julian. In return for that and doing the initial legal work and giving us business advice, we gave him 10% of the company. Ten years later this deal became the model for Y Combinator's. We knew founders needed something like this, because we'd needed it ourselves.\n",
            "\n",
            "At this stage I had a negative net worth, because the thousand dollars or so I had in the bank was more than counterbalanced by what I owed the government in taxes. (Had I diligently set aside the proper proportion of the money I'd made consulting for Interleaf? No, I had not.) So although Robert had his graduate student stipend, I needed that seed funding to live on.\n",
            "\n",
            "We originally hoped to launch in September, but we got more ambitious about the software as we worked on it....\n",
            "-----\n",
            "\n",
            "Score: 0.03 - [8]\n",
            "\n",
            "There were three main parts to the software: the editor, which people used to build sites and which I wrote, the shopping cart, which Robert wrote, and the manager, which kept track of orders and statistics, and which Trevor wrote. In its time, the editor was one of the best general-purpose site builders. I kept the code tight and didn't have to integrate with any other software except Robert's and Trevor's, so it was quite fun to work on. If all I'd had to do was work on this software, the next 3 years would have been the easiest of my life. Unfortunately I had to do a lot more, all of it stuff I was worse at than programming, and the next 3 years were instead the most stressful.\n",
            "\n",
            "There were a lot of startups making ecommerce software in the second half of the 90s. We were determined to be the Microsoft Word, not the Interleaf. Which meant being easy to use and inexpensive. It was lucky for us that we were poor, because that caused us to make Viaweb even more inexpensive than we realized. We charged $100 a month for a small store and $300 a month for a big one....\n",
            "-----\n",
            "\n"
          ]
        }
      ],
      "source": [
        "for node in nodes_with_scores:\n",
        "    print(f\"Score: {node.score:.2f} - {node.text}...\\n-----\\n\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-BONiU7JKNXm"
      },
      "source": [
        "## Use in a Query Engine!\n",
        "\n",
        "Now, we can plug our retriever into a query engine to synthesize natural language responses."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fd6cqDoAKNXm"
      },
      "outputs": [],
      "source": [
        "from llama_index.query_engine import RetrieverQueryEngine\n",
        "\n",
        "query_engine = RetrieverQueryEngine.from_args(retriever)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_JO4BrQLKNXm",
        "outputId": "0cdd979c-9e32-48eb-db88-c183466c8975"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Generated queries:\n",
            "1. What were the major events or milestones in the history of Interleafe and Viaweb?\n",
            "2. Can you provide a timeline of the key developments and achievements of Interleafe and Viaweb?\n",
            "3. What were the successes and failures of Interleafe and Viaweb as companies?\n"
          ]
        }
      ],
      "source": [
        "response = query_engine.query(\"What happened at Interleafe and Viaweb?\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6czVaqDLKNXm",
        "outputId": "fd546163-a71b-4f30-ac76-665426e54ff7"
      },
      "outputs": [
        {
          "data": {
            "text/markdown": [
              "**`Final Response:`** At Interleaf, the author had worked as a consultant and had made some money. However, they did not set aside the proper proportion of the money to pay taxes, resulting in a negative net worth. \n",
              "\n",
              "At Viaweb, the author and their team started a new company that developed software for building websites and managing online stores. They received $10,000 in seed funding from Julian, who was the husband of Idelle. In return for the funding and business advice, Julian received a 10% stake in the company. The software developed by Viaweb was designed to be easy to use and inexpensive, with prices ranging from $100 to $300 per month."
            ],
            "text/plain": [
              "<IPython.core.display.Markdown object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "from llama_index.response.notebook_utils import display_response\n",
        "\n",
        "display_response(response)"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "llama-index-4a-wkI5X-py3.11",
      "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"
    },
    "colab": {
      "provenance": []
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}