RAG/07_Self_Query_RAG/Self_Query_RAG.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "id": "76534cce",
      "metadata": {
        "id": "76534cce"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/jerryjliu/llama_index/blob/main/docs/docs/examples/metadata_extraction/MetadataExtractionSEC.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "b07531d9-7473-480d-bee6-c1ee4cbc207c",
      "metadata": {
        "id": "b07531d9-7473-480d-bee6-c1ee4cbc207c"
      },
      "source": [
        "# Extracting Metadata for Better Document Indexing and Understanding\n",
        "\n",
        "In many cases, especially with long documents, a chunk of text may lack the context necessary to disambiguate the chunk from other similar chunks of text. One method of addressing this is manually labelling each chunk in our dataset or knowledge base. However, this can be labour intensive and time consuming for a large number or continually updated set of documents.\n",
        "\n",
        "To combat this, we use LLMs to extract certain contextual information relevant to the document to better help the retrieval and language models disambiguate similar-looking passages.\n",
        "\n",
        "We do this through our brand-new `Metadata Extractor` modules."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "f887966e",
      "metadata": {
        "id": "f887966e"
      },
      "source": [
        "If you're opening this Notebook on colab, you will probably need to install LlamaIndex 🦙."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "1a2eaf14",
      "metadata": {
        "id": "1a2eaf14"
      },
      "outputs": [],
      "source": [
        "%pip install llama-index-llms-openai\n",
        "%pip install llama-index-extractors-entity"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "69d1b1f2",
      "metadata": {
        "id": "69d1b1f2"
      },
      "outputs": [],
      "source": [
        "!pip install llama-index"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "91c807cc-1334-4f92-8a9e-9ccd702f3578",
      "metadata": {
        "id": "91c807cc-1334-4f92-8a9e-9ccd702f3578"
      },
      "outputs": [],
      "source": [
        "import nest_asyncio\n",
        "\n",
        "nest_asyncio.apply()\n",
        "\n",
        "import os\n",
        "import openai\n",
        "\n",
        "os.environ[\"OPENAI_API_KEY\"] = \"YOUR_API_KEY_HERE\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "0adb8e4a-6728-4073-8256-8b3be4ab1e64",
      "metadata": {
        "id": "0adb8e4a-6728-4073-8256-8b3be4ab1e64"
      },
      "outputs": [],
      "source": [
        "from llama_index.llms.openai import OpenAI\n",
        "from llama_index.core.schema import MetadataMode"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "a0231dff-7443-46bf-9b9d-759198d3408e",
      "metadata": {
        "id": "a0231dff-7443-46bf-9b9d-759198d3408e"
      },
      "outputs": [],
      "source": [
        "llm = OpenAI(temperature=0.1, model=\"gpt-3.5-turbo\", max_tokens=512)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "2db2cf90-f295-4a3d-a47c-4b2b1dd2d7c5",
      "metadata": {
        "id": "2db2cf90-f295-4a3d-a47c-4b2b1dd2d7c5"
      },
      "source": [
        "We create a node parser that extracts the document title and hypothetical question embeddings relevant to the document chunk.\n",
        "\n",
        "We also show how to instantiate the `SummaryExtractor` and `KeywordExtractor`, as well as how to create your own custom extractor\n",
        "based on the `BaseExtractor` base class"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "3bda151d-6fb8-427e-82fc-0f3bb469d705",
      "metadata": {
        "id": "3bda151d-6fb8-427e-82fc-0f3bb469d705"
      },
      "outputs": [],
      "source": [
        "from llama_index.core.extractors import (\n",
        "    SummaryExtractor,\n",
        "    QuestionsAnsweredExtractor,\n",
        "    TitleExtractor,\n",
        "    KeywordExtractor,\n",
        "    BaseExtractor,\n",
        ")\n",
        "from llama_index.extractors.entity import EntityExtractor\n",
        "from llama_index.core.node_parser import TokenTextSplitter\n",
        "\n",
        "text_splitter = TokenTextSplitter(\n",
        "    separator=\" \", chunk_size=512, chunk_overlap=128\n",
        ")\n",
        "\n",
        "\n",
        "class CustomExtractor(BaseExtractor):\n",
        "    def extract(self, nodes):\n",
        "        metadata_list = [\n",
        "            {\n",
        "                \"custom\": (\n",
        "                    node.metadata[\"document_title\"]\n",
        "                    + \"\\n\"\n",
        "                    + node.metadata[\"excerpt_keywords\"]\n",
        "                )\n",
        "            }\n",
        "            for node in nodes\n",
        "        ]\n",
        "        return metadata_list\n",
        "\n",
        "\n",
        "extractors = [\n",
        "    TitleExtractor(nodes=5, llm=llm),\n",
        "    QuestionsAnsweredExtractor(questions=3, llm=llm),\n",
        "    # EntityExtractor(prediction_threshold=0.5),\n",
        "    # SummaryExtractor(summaries=[\"prev\", \"self\"], llm=llm),\n",
        "    # KeywordExtractor(keywords=10, llm=llm),\n",
        "    # CustomExtractor()\n",
        "]\n",
        "\n",
        "transformations = [text_splitter] + extractors"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "c72c45a9-dcad-4925-b2f7-d25fe5d80c2d",
      "metadata": {
        "id": "c72c45a9-dcad-4925-b2f7-d25fe5d80c2d"
      },
      "outputs": [],
      "source": [
        "from llama_index.core import SimpleDirectoryReader"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "e4e54937-e9e7-48ed-8600-72cd2f3c529b",
      "metadata": {
        "id": "e4e54937-e9e7-48ed-8600-72cd2f3c529b"
      },
      "source": [
        "We first load the 10k annual SEC report for Uber and Lyft for the years 2019 and 2020 respectively."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "a2e5ef50-82ef-4936-bbc2-c022f67007a0",
      "metadata": {
        "id": "a2e5ef50-82ef-4936-bbc2-c022f67007a0"
      },
      "outputs": [],
      "source": [
        "!mkdir -p data\n",
        "!wget -O \"data/10k-132.pdf\" \"https://www.dropbox.com/scl/fi/6dlqdk6e2k1mjhi8dee5j/uber.pdf?rlkey=2jyoe49bg2vwdlz30l76czq6g&dl=1\"\n",
        "!wget -O \"data/10k-vFinal.pdf\" \"https://www.dropbox.com/scl/fi/qn7g3vrk5mqb18ko4e5in/lyft.pdf?rlkey=j6jxtjwo8zbstdo4wz3ns8zoj&dl=1\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "38a46bf6-9539-4ac2-ad97-eb909992b94d",
      "metadata": {
        "id": "38a46bf6-9539-4ac2-ad97-eb909992b94d"
      },
      "outputs": [],
      "source": [
        "# Note the uninformative document file name, which may be a common scenario in a production setting\n",
        "uber_docs = SimpleDirectoryReader(input_files=[\"data/10k-132.pdf\"]).load_data()\n",
        "uber_front_pages = uber_docs[0:3]\n",
        "uber_content = uber_docs[63:69]\n",
        "uber_docs = uber_front_pages + uber_content"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "269f8ecc-489d-435f-9d81-a9c64fd4d400",
      "metadata": {
        "id": "269f8ecc-489d-435f-9d81-a9c64fd4d400"
      },
      "outputs": [],
      "source": [
        "from llama_index.core.ingestion import IngestionPipeline\n",
        "\n",
        "pipeline = IngestionPipeline(transformations=transformations)\n",
        "\n",
        "uber_nodes = pipeline.run(documents=uber_docs)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "8da4d824-d518-4d37-8322-a35adac05157",
      "metadata": {
        "id": "8da4d824-d518-4d37-8322-a35adac05157",
        "outputId": "4a4d059b-da9d-43b6-860e-6eed04661ea4"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "{'page_label': '2',\n",
              " 'file_name': '10k-132.pdf',\n",
              " 'document_title': 'Exploring the Diverse Landscape of 2019: A Comprehensive Annual Report on Uber Technologies, Inc.',\n",
              " 'questions_this_excerpt_can_answer': '1. How many countries does Uber operate in?\\n2. What is the total gross bookings of Uber in 2019?\\n3. How many trips did Uber facilitate in 2019?'}"
            ]
          },
          "execution_count": null,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "uber_nodes[1].metadata"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "93e70bfb-6c02-401b-be91-3827f358b22c",
      "metadata": {
        "id": "93e70bfb-6c02-401b-be91-3827f358b22c"
      },
      "outputs": [],
      "source": [
        "# Note the uninformative document file name, which may be a common scenario in a production setting\n",
        "lyft_docs = SimpleDirectoryReader(\n",
        "    input_files=[\"data/10k-vFinal.pdf\"]\n",
        ").load_data()\n",
        "lyft_front_pages = lyft_docs[0:3]\n",
        "lyft_content = lyft_docs[68:73]\n",
        "lyft_docs = lyft_front_pages + lyft_content"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "e3720b40-c50c-4185-aaf4-289ff8ab057e",
      "metadata": {
        "id": "e3720b40-c50c-4185-aaf4-289ff8ab057e"
      },
      "outputs": [],
      "source": [
        "from llama_index.core.ingestion import IngestionPipeline\n",
        "\n",
        "pipeline = IngestionPipeline(transformations=transformations)\n",
        "\n",
        "lyft_nodes = pipeline.run(documents=lyft_docs)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "98740f96-afdd-45ff-bcc0-2c50965a7349",
      "metadata": {
        "id": "98740f96-afdd-45ff-bcc0-2c50965a7349",
        "outputId": "ef660ef6-0b55-4d97-bd20-ecdcabd2990e"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "{'page_label': '2',\n",
              " 'file_name': '10k-vFinal.pdf',\n",
              " 'document_title': 'Lyft, Inc. Annual Report on Form 10-K for the Fiscal Year Ended December 31, 2020',\n",
              " 'questions_this_excerpt_can_answer': \"1. Has Lyft, Inc. filed a report on and attestation to its management's assessment of the effectiveness of its internal control over financial reporting under Section 404(b) of the Sarbanes-Oxley Act?\\n2. Is Lyft, Inc. considered a shell company according to Rule 12b-2 of the Exchange Act?\\n3. What was the aggregate market value of Lyft, Inc.'s common stock held by non-affiliates on June 30, 2020?\"}"
            ]
          },
          "execution_count": null,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "lyft_nodes[2].metadata"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "ddd5805f-c459-40ae-a21c-5fa0de750a60",
      "metadata": {
        "id": "ddd5805f-c459-40ae-a21c-5fa0de750a60"
      },
      "source": [
        "Since we are asking fairly sophisticated questions, we utilize a subquestion query engine for all QnA pipelines below, and prompt it to pay more attention to the relevance of the retrieved sources."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "302bb085-86cc-4b76-a452-67bc826b292d",
      "metadata": {
        "id": "302bb085-86cc-4b76-a452-67bc826b292d"
      },
      "outputs": [],
      "source": [
        "from llama_index.core.question_gen import LLMQuestionGenerator\n",
        "from llama_index.core.question_gen.prompts import (\n",
        "    DEFAULT_SUB_QUESTION_PROMPT_TMPL,\n",
        ")\n",
        "\n",
        "\n",
        "question_gen = LLMQuestionGenerator.from_defaults(\n",
        "    llm=llm,\n",
        "    prompt_template_str=\"\"\"\n",
        "        Follow the example, but instead of giving a question, always prefix the question\n",
        "        with: 'By first identifying and quoting the most relevant sources, '.\n",
        "        \"\"\"\n",
        "    + DEFAULT_SUB_QUESTION_PROMPT_TMPL,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "12cd959f-80ac-4d10-9246-d27cc6c1096a",
      "metadata": {
        "id": "12cd959f-80ac-4d10-9246-d27cc6c1096a"
      },
      "source": [
        "## Querying an Index With No Extra Metadata"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "37dd8992-3716-44da-9309-154fb5946e98",
      "metadata": {
        "id": "37dd8992-3716-44da-9309-154fb5946e98",
        "outputId": "645376e2-fbd5-40b3-817f-4cc337951063"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "LLM sees:\n",
            " [Excerpt from document]\n",
            "page_label: 65\n",
            "file_name: 10k-132.pdf\n",
            "Excerpt:\n",
            "-----\n",
            "See the section titled “Reconciliations of Non-GAAP Financial Measures” for our definition and a \n",
            "reconciliation of net income (loss) attributable to  Uber Technologies, Inc. to Adjusted EBITDA. \n",
            "            \n",
            "  Year Ended December 31,   2017 to 2018   2018 to 2019   \n",
            "(In millions, exce pt percenta ges)  2017   2018   2019   % Chan ge  % Chan ge  \n",
            "Adjusted EBITDA ................................  $ (2,642) $ (1,847) $ (2,725)  30%  (48)%\n",
            "-----\n"
          ]
        }
      ],
      "source": [
        "from copy import deepcopy\n",
        "\n",
        "nodes_no_metadata = deepcopy(uber_nodes) + deepcopy(lyft_nodes)\n",
        "for node in nodes_no_metadata:\n",
        "    node.metadata = {\n",
        "        k: node.metadata[k]\n",
        "        for k in node.metadata\n",
        "        if k in [\"page_label\", \"file_name\"]\n",
        "    }\n",
        "print(\n",
        "    \"LLM sees:\\n\",\n",
        "    (nodes_no_metadata)[9].get_content(metadata_mode=MetadataMode.LLM),\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "b8ff619d-67ed-4263-bfc7-2a7a1b7320e7",
      "metadata": {
        "id": "b8ff619d-67ed-4263-bfc7-2a7a1b7320e7"
      },
      "outputs": [],
      "source": [
        "from llama_index.core import VectorStoreIndex\n",
        "from llama_index.core.query_engine import SubQuestionQueryEngine\n",
        "from llama_index.core.tools import QueryEngineTool, ToolMetadata"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "028a65d7-8065-4798-acec-1c3486633e14",
      "metadata": {
        "id": "028a65d7-8065-4798-acec-1c3486633e14"
      },
      "outputs": [],
      "source": [
        "index_no_metadata = VectorStoreIndex(\n",
        "    nodes=nodes_no_metadata,\n",
        ")\n",
        "engine_no_metadata = index_no_metadata.as_query_engine(\n",
        "    similarity_top_k=10, llm=OpenAI(model=\"gpt-4\")\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "73ea9e05-ff5a-49b6-8e52-139d156cde47",
      "metadata": {
        "id": "73ea9e05-ff5a-49b6-8e52-139d156cde47"
      },
      "outputs": [],
      "source": [
        "final_engine_no_metadata = SubQuestionQueryEngine.from_defaults(\n",
        "    query_engine_tools=[\n",
        "        QueryEngineTool(\n",
        "            query_engine=engine_no_metadata,\n",
        "            metadata=ToolMetadata(\n",
        "                name=\"sec_filing_documents\",\n",
        "                description=\"financial information on companies\",\n",
        "            ),\n",
        "        )\n",
        "    ],\n",
        "    question_gen=question_gen,\n",
        "    use_async=True,\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "fd5a3e51-e252-4e24-bc2b-fbc32ce078dd",
      "metadata": {
        "id": "fd5a3e51-e252-4e24-bc2b-fbc32ce078dd",
        "outputId": "e284235f-2529-419e-d4fd-98bc1b6816ea"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Generated 4 sub questions.\n",
            "\u001b[36;1m\u001b[1;3m[sec_filing_documents] Q: What was the cost due to research and development for Uber in 2019\n",
            "\u001b[0m\u001b[33;1m\u001b[1;3m[sec_filing_documents] Q: What was the cost due to sales and marketing for Uber in 2019\n",
            "\u001b[0m\u001b[38;5;200m\u001b[1;3m[sec_filing_documents] Q: What was the cost due to research and development for Lyft in 2019\n",
            "\u001b[0m\u001b[32;1m\u001b[1;3m[sec_filing_documents] Q: What was the cost due to sales and marketing for Lyft in 2019\n",
            "\u001b[0m\u001b[33;1m\u001b[1;3m[sec_filing_documents] A: The cost due to sales and marketing for Uber in 2019 was $814,122 in thousands.\n",
            "\u001b[0m\u001b[36;1m\u001b[1;3m[sec_filing_documents] A: The cost due to research and development for Uber in 2019 was $1,505,640 in thousands.\n",
            "\u001b[0m\u001b[38;5;200m\u001b[1;3m[sec_filing_documents] A: The cost of research and development for Lyft in 2019 was $1,505,640 in thousands.\n",
            "\u001b[0m\u001b[32;1m\u001b[1;3m[sec_filing_documents] A: The cost due to sales and marketing for Lyft in 2019 was $814,122 in thousands.\n",
            "\u001b[0m{\n",
            "  \"Uber\": {\n",
            "    \"Research and Development\": 1505.64,\n",
            "    \"Sales and Marketing\": 814.122\n",
            "  },\n",
            "  \"Lyft\": {\n",
            "    \"Research and Development\": 1505.64,\n",
            "    \"Sales and Marketing\": 814.122\n",
            "  }\n",
            "}\n"
          ]
        }
      ],
      "source": [
        "response_no_metadata = final_engine_no_metadata.query(\n",
        "    \"\"\"\n",
        "    What was the cost due to research and development v.s. sales and marketing for uber and lyft in 2019 in millions of USD?\n",
        "    Give your answer as a JSON.\n",
        "    \"\"\"\n",
        ")\n",
        "print(response_no_metadata.response)\n",
        "# Correct answer:\n",
        "# {\"Uber\": {\"Research and Development\": 4836, \"Sales and Marketing\": 4626},\n",
        "#  \"Lyft\": {\"Research and Development\": 1505.6, \"Sales and Marketing\": 814 }}"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "e9dafdad-c18c-4e0f-8a35-b691ca73e1f2",
      "metadata": {
        "id": "e9dafdad-c18c-4e0f-8a35-b691ca73e1f2"
      },
      "source": [
        "**RESULT**: As we can see, the QnA agent does not seem to know where to look for the right documents. As a result it gets the Lyft and Uber data completely mixed up."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "9878905f-c47d-46b2-9ad9-063538e717e1",
      "metadata": {
        "id": "9878905f-c47d-46b2-9ad9-063538e717e1"
      },
      "source": [
        "## Querying an Index With Extracted Metadata"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "97f00a18-e9e6-47db-bef5-cbf5bb5016be",
      "metadata": {
        "id": "97f00a18-e9e6-47db-bef5-cbf5bb5016be",
        "outputId": "8af76fff-35cf-4898-de02-3770c8b1b66f"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "LLM sees:\n",
            " [Excerpt from document]\n",
            "page_label: 65\n",
            "file_name: 10k-132.pdf\n",
            "document_title: Exploring the Diverse Landscape of 2019: A Comprehensive Annual Report on Uber Technologies, Inc.\n",
            "Excerpt:\n",
            "-----\n",
            "See the section titled “Reconciliations of Non-GAAP Financial Measures” for our definition and a \n",
            "reconciliation of net income (loss) attributable to  Uber Technologies, Inc. to Adjusted EBITDA. \n",
            "            \n",
            "  Year Ended December 31,   2017 to 2018   2018 to 2019   \n",
            "(In millions, exce pt percenta ges)  2017   2018   2019   % Chan ge  % Chan ge  \n",
            "Adjusted EBITDA ................................  $ (2,642) $ (1,847) $ (2,725)  30%  (48)%\n",
            "-----\n"
          ]
        }
      ],
      "source": [
        "print(\n",
        "    \"LLM sees:\\n\",\n",
        "    (uber_nodes + lyft_nodes)[9].get_content(metadata_mode=MetadataMode.LLM),\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "c7d255de-3034-4035-93bc-45d535ce1700",
      "metadata": {
        "id": "c7d255de-3034-4035-93bc-45d535ce1700"
      },
      "outputs": [],
      "source": [
        "index = VectorStoreIndex(\n",
        "    nodes=uber_nodes + lyft_nodes,\n",
        ")\n",
        "engine = index.as_query_engine(similarity_top_k=10, llm=OpenAI(model=\"gpt-4\"))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "bbe42516-a2ca-4986-9012-cb15682323f5",
      "metadata": {
        "id": "bbe42516-a2ca-4986-9012-cb15682323f5"
      },
      "outputs": [],
      "source": [
        "final_engine = SubQuestionQueryEngine.from_defaults(\n",
        "    query_engine_tools=[\n",
        "        QueryEngineTool(\n",
        "            query_engine=engine,\n",
        "            metadata=ToolMetadata(\n",
        "                name=\"sec_filing_documents\",\n",
        "                description=\"financial information on companies.\",\n",
        "            ),\n",
        "        )\n",
        "    ],\n",
        "    question_gen=question_gen,\n",
        "    use_async=True,\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "f48ac2d9-58e9-4b98-9bad-b8ce1eea7934",
      "metadata": {
        "id": "f48ac2d9-58e9-4b98-9bad-b8ce1eea7934",
        "outputId": "4cfdb262-3afb-4de6-d17f-31fbf38869c9"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Generated 4 sub questions.\n",
            "\u001b[36;1m\u001b[1;3m[sec_filing_documents] Q: What was the cost due to research and development for Uber in 2019\n",
            "\u001b[0m\u001b[33;1m\u001b[1;3m[sec_filing_documents] Q: What was the cost due to sales and marketing for Uber in 2019\n",
            "\u001b[0m\u001b[38;5;200m\u001b[1;3m[sec_filing_documents] Q: What was the cost due to research and development for Lyft in 2019\n",
            "\u001b[0m\u001b[32;1m\u001b[1;3m[sec_filing_documents] Q: What was the cost due to sales and marketing for Lyft in 2019\n",
            "\u001b[0m\u001b[33;1m\u001b[1;3m[sec_filing_documents] A: The cost due to sales and marketing for Uber in 2019 was $4,626 million.\n",
            "\u001b[0m\u001b[36;1m\u001b[1;3m[sec_filing_documents] A: The cost due to research and development for Uber in 2019 was $4,836 million.\n",
            "\u001b[0m\u001b[32;1m\u001b[1;3m[sec_filing_documents] A: The cost due to sales and marketing for Lyft in 2019 was $814,122 in thousands.\n",
            "\u001b[0m\u001b[38;5;200m\u001b[1;3m[sec_filing_documents] A: The cost of research and development for Lyft in 2019 was $1,505,640 in thousands.\n",
            "\u001b[0m{\n",
            "  \"Uber\": {\n",
            "    \"Research and Development\": 4836,\n",
            "    \"Sales and Marketing\": 4626\n",
            "  },\n",
            "  \"Lyft\": {\n",
            "    \"Research and Development\": 1505.64,\n",
            "    \"Sales and Marketing\": 814.122\n",
            "  }\n",
            "}\n"
          ]
        }
      ],
      "source": [
        "response = final_engine.query(\n",
        "    \"\"\"\n",
        "    What was the cost due to research and development v.s. sales and marketing for uber and lyft in 2019 in millions of USD?\n",
        "    Give your answer as a JSON.\n",
        "    \"\"\"\n",
        ")\n",
        "print(response.response)\n",
        "# Correct answer:\n",
        "# {\"Uber\": {\"Research and Development\": 4836, \"Sales and Marketing\": 4626},\n",
        "#  \"Lyft\": {\"Research and Development\": 1505.6, \"Sales and Marketing\": 814 }}"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "8bee6d91-84f4-4bde-89dc-d010f9aebc3e",
      "metadata": {
        "id": "8bee6d91-84f4-4bde-89dc-d010f9aebc3e"
      },
      "source": [
        "**RESULT**: As we can see, the LLM answers the questions correctly."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "14826bae-4032-4886-87a8-50f9f28d7ace",
      "metadata": {
        "id": "14826bae-4032-4886-87a8-50f9f28d7ace"
      },
      "source": [
        "### Challenges Identified in the Problem Domain\n",
        "\n",
        "In this example, we observed that the search quality as provided by vector embeddings was rather poor. This was likely due to highly dense financial documents that were likely not representative of the training set for the model.\n",
        "\n",
        "In order to improve the search quality, other methods of neural search that employ more keyword-based approaches may help, such as ColBERTv2/PLAID. In particular, this would help in matching on particular keywords to identify high-relevance chunks.\n",
        "\n",
        "Other valid steps may include utilizing models that are fine-tuned on financial datasets such as Bloomberg GPT.\n",
        "\n",
        "Finally, we can help to further enrich the metadata by providing more contextual information regarding the surrounding context that the chunk is located in.\n",
        "\n",
        "### Improvements to this Example\n",
        "Generally, this example can be improved further with more rigorous evaluation of both the metadata extraction accuracy, and the accuracy and recall of the QnA pipeline. Further, incorporating a larger set of documents as well as the full length documents, which may provide more confounding passages that are difficult to disambiguate, could further stresss test the system we have built and suggest further improvements."
      ]
    }
  ],
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "llama-index",
      "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": 5
}