third_party/LlamaIndex/RAG.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "535faa54",
   "metadata": {},
   "source": [
    "<a href=\"https://colab.research.google.com/github/mistralai/cookbook/blob/main/third_party/LlamaIndex/RAG.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b801d333-af2e-4a83-b74c-6c62414ca4be",
   "metadata": {},
   "source": [
    "## Installation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a2efee20-2cc8-494f-9167-c88dfbac7070",
   "metadata": {},
   "source": [
    "# RAG Pipeline with LlamaIndex\n",
    "\n",
    "In this notebook we will look into building RAG with LlamaIndex using `MistralAI LLM and Embedding Model`. Additionally, we will look into using Index as Retreiver.\n",
    "\n",
    "1. Basic RAG pipeline.\n",
    "2. Index as Retriever."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eeded886-d970-4e19-824f-4602e72869dd",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install llama-index \n",
    "!pip install llama-index-embeddings-mistralai\n",
    "!pip install llama-index-llms-mistralai"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d5d4030c-d4ec-41e5-a013-237cad0dcf38",
   "metadata": {},
   "source": [
    "## Setup API Keys"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "febdddf2-90f1-42f0-966c-56ec6c7a41d9",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ['MISTRAL_API_KEY'] = '<YOUR MISTRALAI API KEY>'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "23c2257e-cebc-4458-8c9c-4092bf71148d",
   "metadata": {},
   "source": [
    "## Basic RAG pipeline\n",
    "\n",
    "Following are the steps involved in Builiding a basic RAG pipeline.\n",
    "\n",
    "1. Setup LLM and Embedding Model\n",
    "2. Download Data\n",
    "3. Load Data\n",
    "4. Create Nodes\n",
    "5. Create Index\n",
    "6. Create Query Engine\n",
    "7. Querying\n",
    "\n",
    "Query Engine combines `Retrieval` and `Response Synthesis` modules to generate response for the given query."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "68a58e1b-45f1-4b6a-926c-3ded4bc45a94",
   "metadata": {},
   "source": [
    "### Setup LLM and Embedding Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "fd577225-c6cd-4826-9cb1-146798834622",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.llms.mistralai import MistralAI\n",
    "from llama_index.embeddings.mistralai import MistralAIEmbedding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "af148c39-b188-4d53-a784-8f8d5ca71c7a",
   "metadata": {},
   "outputs": [],
   "source": [
    "llm = MistralAI(model='mistral-large')\n",
    "embed_model = MistralAIEmbedding()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "4d29d436-eff9-4d07-8062-27efcbe27b98",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core import Settings\n",
    "Settings.llm = llm\n",
    "Settings.embed_model = embed_model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "08cf9770-4fff-4e27-8367-51a5e012b6f3",
   "metadata": {},
   "source": [
    "### Download Data\n",
    "\n",
    "We will use `Uber 2021 10K SEC Filings` for the demonstration."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "2e86e81f-1211-411a-8a6c-9a351bc2164f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--2024-03-29 13:12:05--  https://raw.githubusercontent.com/run-llama/llama_index/main/docs/docs/examples/data/10k/uber_2021.pdf\n",
      "Resolving raw.githubusercontent.com (raw.githubusercontent.com)... 185.199.110.133, 185.199.109.133, 185.199.108.133, ...\n",
      "Connecting to raw.githubusercontent.com (raw.githubusercontent.com)|185.199.110.133|:443... connected.\n",
      "HTTP request sent, awaiting response... 200 OK\n",
      "Length: 1880483 (1.8M) [application/octet-stream]\n",
      "Saving to: ‘./uber_2021.pdf’\n",
      "\n",
      "./uber_2021.pdf     100%[===================>]   1.79M  --.-KB/s    in 0.1s    \n",
      "\n",
      "2024-03-29 13:12:05 (13.7 MB/s) - ‘./uber_2021.pdf’ saved [1880483/1880483]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/docs/examples/data/10k/uber_2021.pdf' -O './uber_2021.pdf'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7c7e92e8-6914-4fb3-b9ac-b95302e1f7ee",
   "metadata": {},
   "source": [
    "### Load Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "3d0c6234-4a2f-441c-941d-5abde32053ce",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core import SimpleDirectoryReader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "f9fdc590-6bd0-4625-a4cd-ba8d89de4120",
   "metadata": {},
   "outputs": [],
   "source": [
    "documents = SimpleDirectoryReader(input_files=[\"./uber_2021.pdf\"]).load_data()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fedd733a-b238-4f4d-a19c-c579b7056d3d",
   "metadata": {},
   "source": [
    "### Create Nodes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "24b0a0e6-1f3a-44dd-9aca-0d90eb77a4d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.node_parser import TokenTextSplitter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "7457510e-f94b-4183-9fa6-c66c055c9db9",
   "metadata": {},
   "outputs": [],
   "source": [
    "splitter = TokenTextSplitter(\n",
    "    chunk_size=512,\n",
    "    chunk_overlap=0,\n",
    ")\n",
    "\n",
    "nodes = splitter.get_nodes_from_documents(documents)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "594d7b13-afb6-487c-8c55-091068ec3fd0",
   "metadata": {},
   "source": [
    "### Create Index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "b76b2474-90ab-4d79-9f5a-9a00be04e366",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core import VectorStoreIndex"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "3683029c-04e2-46ab-9ae0-af21bb57c5e9",
   "metadata": {},
   "outputs": [],
   "source": [
    "index = VectorStoreIndex(nodes)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cf96242e-bd2c-433e-9103-dbc2ecfaa2be",
   "metadata": {},
   "source": [
    "### Create Query Engine"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "ae53a581-a5e1-4d7c-b959-7295a75293f1",
   "metadata": {},
   "outputs": [],
   "source": [
    "query_engine = index.as_query_engine(similarity_top_k=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "824b8af1-bdb0-4a35-9ba0-c5cd05885c9e",
   "metadata": {},
   "source": [
    "### Querying"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "35ead86e-e43b-4405-ba5b-dc4366df9976",
   "metadata": {},
   "outputs": [],
   "source": [
    "response = query_engine.query(\"What is the revenue of Uber in 2021?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "ce1ef94e-af38-45b4-b8f8-e94c0c21b112",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The total revenue for Uber in 2021 was $17,455 million. This includes revenue from various offerings such as Mobility, Delivery, Freight, and All Other revenue streams. The Mobility revenue was $6,953 million, Delivery revenue was $8,362 million, Freight revenue was $2,132 million, and All Other revenue was $8 million.\n"
     ]
    }
   ],
   "source": [
    "print(response)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ba74fa8a-1794-4b8a-9fe0-3c6205a6d190",
   "metadata": {},
   "source": [
    "## Index as Retriever\n",
    "\n",
    "We can make use of created index as a `Retriever`. Retriever helps you to retrieve relevant chunks/ nodes for the given user query."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ba0becc3-2ac6-457a-8653-c59934f61a27",
   "metadata": {},
   "source": [
    "### Create Retriever"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "8017fec5-80ed-45b2-81f0-f8e111ab6197",
   "metadata": {},
   "outputs": [],
   "source": [
    "retriever = index.as_retriever(similarity_top_k = 2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fcdd70f4-3db2-4775-a5e6-50b48ecfd200",
   "metadata": {},
   "source": [
    "### Retrieve relevant nodes for a Query"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "a8c3a7b1-e0eb-4c08-bffd-134aae1d1d0f",
   "metadata": {},
   "outputs": [],
   "source": [
    "retrieved_nodes = retriever.retrieve(\"What is the revenue of Uber in 2021?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "e2f0bdd4-4b3c-42ac-91b3-54eb4b4f3ce5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "**Node ID:** 96264fe5-bc88-4cf8-8905-a9691c39a5c9<br>**Similarity:** 0.8679960824403077<br>**Text:** Note 2 – RevenueThe\n",
       " following  tables  present  our  revenues  disaggregated  by  offering  and  geographical  region.  Revenue  by  geographical  region  is  based  on  where  thetransaction\n",
       " occurred. This level of disaggregation takes into consideration how the nature, amount, timing, and uncertainty of revenue and cash flows are affectedby economic factors. Revenue \n",
       "is presented in the following tables for the years ended December 31, 2019, 2020 and 2021, respectively (in millions):Year Ended December 31,\n",
       "2019\n",
       "2020 2021 Mobility revenue \n",
       "$ 10,707 $ 6,089 $ 6,953 Delivery revenue \n",
       "1,401 3,904 8,362 Freight revenue\n",
       "731 1,011 2,132 All Other revenue\n",
       "161 135 8 Total revenue\n",
       "$ 13,000 $ 11,139 $ 17,455  We\n",
       " offer subscription memberships to end-users including Uber One, Uber Pass, Rides Pass, and Eats Pass (“Subscription”). We recognize Subscriptionfees\n",
       " ratably over the life of the pass. We allocate Subscription fees earned to Mobility and Delivery revenue on a proportional basis, b...<br>"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/markdown": [
       "**Node ID:** 653f0be9-ecfc-4fac-9488-afbd65f44ef2<br>**Similarity:** 0.8590901940430307<br>**Text:** COVID-19.\n",
       "Revenue\n",
       " was $17.5 billion, or up 57% year-over-year, reflecting the overall growth in our Delivery business and an increase in Freight revenue attributable tothe\n",
       " acquisition of Transplace in the fourth quarter of 2021 as well as growth in the number of shippers and carriers on the network combined with an increase involumes with our top shippers.\n",
       "Net\n",
       " loss attributable to Uber Technologies, Inc. was $496 million, a 93% improvement year-over-year, driven by a $1.6 billion pre-tax gain on the sale of ourATG\n",
       " Business to Aurora, a $1.6 billion pre-tax  net benefit relating to Uber’s equity investments, as  well as reductions in our fixed cost structure and increasedvariable cost effi\n",
       "ciencies. Net loss attributable to Uber Technologies, Inc. also included $1.2 billion of stock-based compensation expense.Adjusted\n",
       " EBITDA loss was $774 million, improving $1.8 billion from 2020 with Mobility Adjusted EBITDA profit of $1.6 billion. Additionally, DeliveryAdjusted\n",
       " EBITDA loss of...<br>"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from llama_index.core.response.notebook_utils import display_source_node\n",
    "\n",
    "for node in retrieved_nodes:\n",
    "    display_source_node(node, source_length=1000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ac4cd64d-a800-4d60-b2d6-c529fd8c8cc8",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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",
   "version": "3.11.3 (main, Apr  7 2023, 19:08:44) [Clang 13.0.0 (clang-1300.0.29.30)]"
  },
  "vscode": {
   "interpreter": {
    "hash": "b0fa6594d8f4cbf19f97940f81e996739fb7646882a419484c72d19e05852a7e"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}