Hybrid_Search.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "authorship_tag": "ABX9TyMUQHRTuOH5w9vuDJMk/ZyW",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/Anil-matcha/LlamaIndex-tutorials/blob/main/Hybrid_Search.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Custom Retriever with Hybrid Search\n",
        "\n",
        "Keyword based search was the initial form of search used in information retrieval systems. Then recently we have Vector db based search which works based on semantic similarity.\n",
        "\n",
        "It is not always necessary that a Vector db backed search performs better than a keyword based search on a particular query. It can be vice-versa.\n",
        "\n",
        "Thus to overcome this, we can use Hybrid search which results in best of both worlds. Let's discuss how we can achieve this with the help of a Custom Retreiver in LlamaIndex"
      ],
      "metadata": {
        "id": "Gb-Wn7kz5q8_"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BsTd1h5d5cdJ"
      },
      "outputs": [],
      "source": [
        "!pip install llama-index"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import logging\n",
        "import sys\n",
        "\n",
        "logging.basicConfig(stream=sys.stdout, level=logging.INFO)\n",
        "logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))\n",
        "\n",
        "from llama_index import (\n",
        "    VectorStoreIndex,\n",
        "    SimpleKeywordTableIndex,\n",
        "    SimpleDirectoryReader,\n",
        "    ServiceContext,\n",
        "    StorageContext,\n",
        ")\n",
        "from IPython.display import Markdown, display"
      ],
      "metadata": {
        "id": "uM4iEm7S5yD5"
      },
      "execution_count": 2,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "!mkdir data\n",
        "!wget https://github.com/jerryjliu/llama_index/blob/main/examples/paul_graham_essay/data/paul_graham_essay.txt\n",
        "!mv paul_graham_essay.txt data/\n",
        "documents = SimpleDirectoryReader(\"./data/\").load_data()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "FDT5zfhp5yYd",
        "outputId": "622e478f-4f63-445e-b47d-2fc38cf6a67b"
      },
      "execution_count": 4,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "--2023-07-22 18:40:50--  https://github.com/jerryjliu/llama_index/blob/main/examples/paul_graham_essay/data/paul_graham_essay.txt\n",
            "Resolving github.com (github.com)... 192.30.255.113\n",
            "Connecting to github.com (github.com)|192.30.255.113|:443... connected.\n",
            "HTTP request sent, awaiting response... 200 OK\n",
            "Length: 84954 (83K) [text/plain]\n",
            "Saving to: ‘paul_graham_essay.txt’\n",
            "\n",
            "\rpaul_graham_essay.t   0%[                    ]       0  --.-KB/s               \rpaul_graham_essay.t 100%[===================>]  82.96K  --.-KB/s    in 0.02s   \n",
            "\n",
            "2023-07-22 18:40:51 (4.76 MB/s) - ‘paul_graham_essay.txt’ saved [84954/84954]\n",
            "\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "service_context = ServiceContext.from_defaults(chunk_size=1024)\n",
        "node_parser = service_context.node_parser\n",
        "nodes = node_parser.get_nodes_from_documents(documents)\n",
        "storage_context = StorageContext.from_defaults()\n",
        "storage_context.docstore.add_documents(nodes)"
      ],
      "metadata": {
        "id": "Lfktds3D50LQ"
      },
      "execution_count": 5,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import openai\n",
        "openai.api_key = \"your-openai-key\""
      ],
      "metadata": {
        "id": "qJBFXwjc6IMP"
      },
      "execution_count": 7,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "vector_index = VectorStoreIndex(nodes, storage_context=storage_context)\n",
        "keyword_index = SimpleKeywordTableIndex(nodes, storage_context=storage_context)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "98cSPPBW55IY",
        "outputId": "d84a4795-0013-443e-838f-f096dddc83bd"
      },
      "execution_count": 8,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "[nltk_data] Downloading package stopwords to /root/nltk_data...\n",
            "[nltk_data]   Unzipping corpora/stopwords.zip.\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "from llama_index import QueryBundle\n",
        "\n",
        "# import NodeWithScore\n",
        "from llama_index.schema import NodeWithScore\n",
        "\n",
        "# Retrievers\n",
        "from llama_index.retrievers import (\n",
        "    BaseRetriever,\n",
        "    VectorIndexRetriever,\n",
        "    KeywordTableSimpleRetriever,\n",
        ")\n",
        "\n",
        "from typing import List\n",
        "\n",
        "class CustomRetriever(BaseRetriever):\n",
        "    \"\"\"Custom retriever that performs both semantic search and hybrid search.\"\"\"\n",
        "\n",
        "    def __init__(\n",
        "        self,\n",
        "        vector_retriever: VectorIndexRetriever,\n",
        "        keyword_retriever: KeywordTableSimpleRetriever,\n",
        "        mode: str = \"AND\",\n",
        "    ) -> None:\n",
        "        \"\"\"Init params.\"\"\"\n",
        "\n",
        "        self._vector_retriever = vector_retriever\n",
        "        self._keyword_retriever = keyword_retriever\n",
        "        if mode not in (\"AND\", \"OR\"):\n",
        "            raise ValueError(\"Invalid mode.\")\n",
        "        self._mode = mode\n",
        "\n",
        "    def _retrieve(self, query_bundle: QueryBundle) -> List[NodeWithScore]:\n",
        "        \"\"\"Retrieve nodes given query.\"\"\"\n",
        "\n",
        "        vector_nodes = self._vector_retriever.retrieve(query_bundle)\n",
        "        keyword_nodes = self._keyword_retriever.retrieve(query_bundle)\n",
        "\n",
        "        vector_ids = {n.node.node_id for n in vector_nodes}\n",
        "        keyword_ids = {n.node.node_id for n in keyword_nodes}\n",
        "\n",
        "        combined_dict = {n.node.node_id: n for n in vector_nodes}\n",
        "        combined_dict.update({n.node.node_id: n for n in keyword_nodes})\n",
        "\n",
        "        if self._mode == \"AND\":\n",
        "            retrieve_ids = vector_ids.intersection(keyword_ids)\n",
        "        else:\n",
        "            retrieve_ids = vector_ids.union(keyword_ids)\n",
        "\n",
        "        retrieve_nodes = [combined_dict[rid] for rid in retrieve_ids]\n",
        "        return retrieve_nodes"
      ],
      "metadata": {
        "id": "BSzfu6Xw576c"
      },
      "execution_count": 9,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from llama_index import get_response_synthesizer\n",
        "from llama_index.query_engine import RetrieverQueryEngine\n",
        "\n",
        "# define custom retriever\n",
        "vector_retriever = VectorIndexRetriever(index=vector_index, similarity_top_k=2)\n",
        "keyword_retriever = KeywordTableSimpleRetriever(index=keyword_index)\n",
        "custom_retriever = CustomRetriever(vector_retriever, keyword_retriever)\n",
        "\n",
        "# define response synthesizer\n",
        "response_synthesizer = get_response_synthesizer()\n",
        "\n",
        "# assemble query engine\n",
        "custom_query_engine = RetrieverQueryEngine(\n",
        "    retriever=custom_retriever,\n",
        "    response_synthesizer=response_synthesizer,\n",
        ")\n",
        "\n",
        "# vector query engine\n",
        "vector_query_engine = RetrieverQueryEngine(\n",
        "    retriever=vector_retriever,\n",
        "    response_synthesizer=response_synthesizer,\n",
        ")\n",
        "# keyword query engine\n",
        "keyword_query_engine = RetrieverQueryEngine(\n",
        "    retriever=keyword_retriever,\n",
        "    response_synthesizer=response_synthesizer,\n",
        ")"
      ],
      "metadata": {
        "id": "Rv3mhMUO5_62"
      },
      "execution_count": 10,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "response = custom_query_engine.query(\"What did the author do during his time at YC?\")"
      ],
      "metadata": {
        "id": "JwQQ9EkT6B0c"
      },
      "execution_count": 11,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "print(response)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "cjSrPcpa6DW0",
        "outputId": "7abce0ba-74ac-4a9b-c6b5-1b7b38638692"
      },
      "execution_count": 12,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "The author worked on YC, writing essays, developing internal software in Arc, and creating Hacker News. He also helped select and support founders, resolve disputes between cofounders, and fight with people who maltreated the startups. He worked hard, even at the parts he didn't like, and eventually handed YC over to someone else. After his mother's death, he checked out of YC and decided to pursue painting.\n"
          ]
        }
      ]
    }
  ]
}