LlamaIndex_QA_+_Summary.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "authorship_tag": "ABX9TyMUcGEpWYrW1phg4PRIVNjW",
      "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/LlamaIndex_QA_%2B_Summary.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install llama-index"
      ],
      "metadata": {
        "id": "2vPspYq2pXcd"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "jkeqbN1MorPp"
      },
      "outputs": [],
      "source": [
        "import nest_asyncio\n",
        "nest_asyncio.apply()"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "from llama_index import (\n",
        "    VectorStoreIndex,\n",
        "    ListIndex,\n",
        "    SimpleDirectoryReader,\n",
        "    ServiceContext,\n",
        "    StorageContext)"
      ],
      "metadata": {
        "id": "krd8IAm0pL6G"
      },
      "execution_count": 4,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "!wget https://raw.githubusercontent.com/hwchase17/chat-your-data/master/state_of_the_union.txt\n",
        "!mkdir data\n",
        "!mv state_of_the_union.txt data/"
      ],
      "metadata": {
        "id": "tW-R_LqApiKE"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "documents = SimpleDirectoryReader('./data').load_data()"
      ],
      "metadata": {
        "id": "fX4ER_FSp99s"
      },
      "execution_count": 12,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import os\n",
        "os.environ[\"OPENAI_API_KEY\"] = \"openai-key\"\n",
        "service_context = ServiceContext.from_defaults(chunk_size=1024)\n",
        "nodes = service_context.node_parser.get_nodes_from_documents(documents)"
      ],
      "metadata": {
        "id": "BnyPxCn6qqy9"
      },
      "execution_count": 13,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "storage_context = StorageContext.from_defaults()\n",
        "storage_context.docstore.add_documents(nodes)"
      ],
      "metadata": {
        "id": "IbOM375AqvSY"
      },
      "execution_count": 14,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "list_index = ListIndex(nodes, storage_context=storage_context)\n",
        "vector_index = VectorStoreIndex(nodes, storage_context=storage_context)"
      ],
      "metadata": {
        "id": "mCMs4gb9rHe3"
      },
      "execution_count": 15,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "list_query_engine = list_index.as_query_engine(\n",
        "    response_mode='tree_summarize',\n",
        "    use_async=True,\n",
        ")\n",
        "vector_query_engine = vector_index.as_query_engine()"
      ],
      "metadata": {
        "id": "DslYNvuCrJq1"
      },
      "execution_count": 16,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from llama_index.tools.query_engine import QueryEngineTool\n",
        "\n",
        "\n",
        "list_tool = QueryEngineTool.from_defaults(\n",
        "    query_engine=list_query_engine,\n",
        "    description='Useful for summarization questions related to Paul Graham eassy on What I Worked On.',\n",
        ")\n",
        "\n",
        "vector_tool = QueryEngineTool.from_defaults(\n",
        "    query_engine=vector_query_engine,\n",
        "    description='Useful for retrieving specific context from Paul Graham essay on What I Worked On.',\n",
        "\n",
        ")"
      ],
      "metadata": {
        "id": "Pgn29m0qrSel"
      },
      "execution_count": 17,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from llama_index.query_engine.router_query_engine import RouterQueryEngine\n",
        "from llama_index.selectors.llm_selectors import LLMSingleSelector\n",
        "\n",
        "\n",
        "query_engine = RouterQueryEngine(\n",
        "    selector=LLMSingleSelector.from_defaults(),\n",
        "    query_engine_tools=[\n",
        "        list_tool,\n",
        "        vector_tool,\n",
        "    ]\n",
        ")"
      ],
      "metadata": {
        "id": "ms5ovM14rXS5"
      },
      "execution_count": 18,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "print(query_engine.query('What is the summary of the document?').response)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "AX4ldlPiraEA",
        "outputId": "e88c14fa-60fa-438f-83ac-02890722a8e9"
      },
      "execution_count": 23,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "In this speech, President Biden calls on Congress to pass legislation to fight inflation, lower the cost of prescription drugs, cut energy costs, end the shutdown of schools and businesses, continue vaccinating the world, and protect communities by investing in crime prevention and community police officers. He also calls for universal background checks and the repeal of the liability shield for gun manufacturers. Finally, he calls for the passage of the Freedom to Vote Act, the John Lewis Voting Rights Act, and the Disclose Act to protect the right to vote.\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "print(query_engine.query('What is NaTO?').response)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "uKsa35JUrcUr",
        "outputId": "fc4a817c-6428-4f7c-ed1f-b352d85e74f9"
      },
      "execution_count": 22,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "NATO is the North Atlantic Treaty Organization, an intergovernmental military alliance between 30 North American and European countries. It was created after World War II to secure peace and stability in Europe.\n"
          ]
        }
      ]
    }
  ]
}