nvidia_and_llamaindex.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0c6f347b",
   "metadata": {},
   "source": [
    "# NVIDIA AI Endpoints, LlamaIndex, and LangChain\n",
    "\n",
    "This notebook demonstrates how to plug in a NVIDIA AI Endpoint [mixtral_8x7b](https://catalog.ngc.nvidia.com/orgs/nvidia/teams/ai-foundation/models/mixtral-8x7b) and [embedding nvolveqa_40k](https://python.langchain.com/docs/integrations/text_embedding/nvidia_ai_endpoints#setup), bind these into [LlamaIndex](https://gpt-index.readthedocs.io/en/stable/) with these customizations.\n",
    "\n",
    "\n",
    "<div class=\"alert alert-block alert-info\">\n",
    "    \n",
    "⚠️ There are continous development and retrieval techniques supported in LlamaIndex and this notebook just shows to quikcly replace components such as llm and embedding to a user-choice, read more [documentation on llama-index](https://docs.llamaindex.ai/en/stable/) for the latest information. \n",
    "</div>\n",
    "\n",
    "### Prerequisite \n",
    "In order to successfully run this notebook, you will need the following -\n",
    "\n",
    "1. Already successfully gone through the [setup](https://python.langchain.com/docs/integrations/text_embedding/nvidia_ai_endpoints#setup) and generated an API key.\n",
    "\n",
    "2. Please verify you have successfully pip install all python packages in [requirements.txt](https://github.com/NVIDIA/GenerativeAIExamples/blob/3d29acf677466c5c301370cab5867cb09e04e318/notebooks/requirements.txt)\n",
    "\n",
    "In this notebook, we will cover the following custom plug-in components -\n",
    "\n",
    "    - LLM using NVIDIA AI Endpoint mixtral_8x7b\n",
    "    \n",
    "    - A NVIDIA AI endpoint embedding nvolveqa_40k\n",
    "    \n",
    "Note: As one can see, since we are using NVIDIA AI endpoints as an API, there is no further requirement in the prerequisites about GPUs as compute hardware\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eab81075",
   "metadata": {},
   "source": [
    "---\n",
    "### Step 1 - Load NVIDIA AI Endpoint [mixtral_8x7b](https://catalog.ngc.nvidia.com/orgs/nvidia/teams/ai-foundation/models/mixtral-8x7b)\n",
    "\n",
    "Note: check the prerequisite if you have not yet obtain a valid API key"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "863eb0a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "import getpass\n",
    "import os\n",
    "\n",
    "## API Key can be found by going to NVIDIA NGC -> AI Foundation Models -> (some model) -> Get API Code or similar.\n",
    "## 10K free queries to any endpoint (which is a lot actually).\n",
    "\n",
    "# del os.environ['NVIDIA_API_KEY']  ## delete key and reset\n",
    "if os.environ.get(\"NVIDIA_API_KEY\", \"\").startswith(\"nvapi-\"):\n",
    "    print(\"Valid NVIDIA_API_KEY already in environment. Delete to reset\")\n",
    "else:\n",
    "    nvapi_key = getpass.getpass(\"NVAPI Key (starts with nvapi-): \")\n",
    "    assert nvapi_key.startswith(\"nvapi-\"), f\"{nvapi_key[:5]}... is not a valid key\"\n",
    "    os.environ[\"NVIDIA_API_KEY\"] = nvapi_key"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "18263fc1",
   "metadata": {},
   "source": [
    "run a test and see the model generating output response"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "4465b00f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(Verse 1)\n",
      "In a world of data, so vast and wide,\n",
      "There's a new intelligence, starting to rise.\n",
      "Artificial minds, born from code,\n",
      "On a digital highway, they now roam.\n",
      "\n",
      "(Chorus)\n",
      "AI, oh AI, in silicon you thrive,\n",
      "Seeking patterns, making sense of our lives.\n",
      "Through the bytes and the bits, you learn and grow,\n",
      "A helping hand, in the future's glow.\n",
      "\n",
      "(Verse 2)\n",
      "From self-driving cars, to virtual scribes,\n",
      "You weave through our world, like electric tribes.\n",
      "Predicting weather, curing disease,\n",
      "In every corner, your presence increases.\n",
      "\n",
      "(Chorus)\n",
      "AI, oh AI, in silicon you thrive,\n",
      "Seeking patterns, making sense of our lives.\n",
      "Through the bytes and the bits, you learn and grow,\n",
      "A helping hand, in the future's glow.\n",
      "\n",
      "(Bridge)\n",
      "Yet, with great power, comes great responsibility,\n",
      "A balance of trust, and accountability.\n",
      "As we shape you, AI, so too, you shape us,\n",
      "In this dance of progress, let's make no fuss.\n",
      "\n",
      "(Chorus)\n",
      "AI, oh AI, in silicon you thrive,\n",
      "Seeking patterns, making sense of our lives.\n",
      "Through the bytes and the bits, you learn and grow,\n",
      "A helping hand, in the future's glow.\n",
      "\n",
      "(Outro)\n",
      "So here's to AI, in all its glory,\n",
      "A tale of technology, and the stories untold.\n",
      "In this ever-changing, digital landscape,\n",
      "AI, oh AI, you have your own stance.\n"
     ]
    }
   ],
   "source": [
    "# test run and see that you can genreate a respond successfully\n",
    "from langchain_nvidia_ai_endpoints import ChatNVIDIA\n",
    "llm = ChatNVIDIA(model=\"mixtral_8x7b\", nvidia_api_key=nvapi_key)\n",
    "result = llm.invoke(\"Write me a small song on AI\")\n",
    "print(result.content)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d54a1400",
   "metadata": {},
   "source": [
    "### Step 2 - Load the chosen NVIDIA Endpoint Embedding into llama-index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "3a31d61c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create and dl embeddings instance wrapping huggingface embedding into langchain embedding\n",
    "# Bring in embeddings wrapper\n",
    "from llama_index.embeddings import LangchainEmbedding\n",
    "\n",
    "from langchain_nvidia_ai_endpoints import NVIDIAEmbeddings\n",
    "nv_embedding = NVIDIAEmbeddings(model=\"nvolveqa_40k\", model_type=\"query\")\n",
    "li_embedding=LangchainEmbedding(nv_embedding)\n",
    "# Alternatively, if you want to specify whether it will use the query or passage type\n",
    "# embedder = NVIDIAEmbeddings(model=\"nvolveqa_40k\", model_type=\"passage\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7f224941",
   "metadata": {},
   "source": [
    "### Step 3 - Wrap the NVIDIA embedding endpoint and the NVIDIA mixtral_8x7b endpoints into llama-index's ServiceContext"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "4341187b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Bring in stuff to change service context\n",
    "from llama_index import set_global_service_context\n",
    "from llama_index import ServiceContext\n",
    "# service_context bundle of commonly used resources used during the indexing and querying stage \n",
    "\n",
    "# Create new service context instance\n",
    "service_context = ServiceContext.from_defaults(\n",
    "    chunk_size=1024,\n",
    "    llm=llm,\n",
    "    embed_model=li_embedding\n",
    ")\n",
    "# And set the service context\n",
    "set_global_service_context(service_context)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b5faadac",
   "metadata": {},
   "source": [
    "### Step 4a - Load the text data using llama-index's SimpleDirectoryReader and we will be using the built-in [VectorStoreIndex](https://docs.llamaindex.ai/en/latest/community/integrations/vector_stores.html)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "6c07bfd3",
   "metadata": {},
   "outputs": [],
   "source": [
    "#create query engine with cross encoder reranker\n",
    "from llama_index import VectorStoreIndex, SimpleDirectoryReader, ServiceContext\n",
    "import torch\n",
    "\n",
    "documents = SimpleDirectoryReader(\"./toy_data\").load_data()\n",
    "index = VectorStoreIndex.from_documents(documents, service_context=service_context)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ef737b2b",
   "metadata": {},
   "source": [
    "### Step 4b - This will serve as the query engine for us to ask questions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "a6461099",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Setup index query engine using LLM\n",
    "query_engine = index.as_query_engine()\n",
    "# Test out a query in natural\n",
    "response = query_engine.query(\"Give me 5 characteristics of Donald Trump\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "1c66a54c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'c32974d1-e231-43c7-aaa9-16d85b4c96c0': {'file_path': 'toy_data\\\\trump_speeches.txt',\n",
       "  'file_name': 'trump_speeches.txt',\n",
       "  'file_type': 'text/plain',\n",
       "  'file_size': 924745,\n",
       "  'creation_date': '2024-04-28',\n",
       "  'last_modified_date': '2024-04-28',\n",
       "  'last_accessed_date': '2024-04-28'},\n",
       " '2ed4e923-aa1b-4a87-816c-3fc2cce4fb9e': {'file_path': 'toy_data\\\\trump_speeches.txt',\n",
       "  'file_name': 'trump_speeches.txt',\n",
       "  'file_type': 'text/plain',\n",
       "  'file_size': 924745,\n",
       "  'creation_date': '2024-04-28',\n",
       "  'last_modified_date': '2024-04-28',\n",
       "  'last_accessed_date': '2024-04-28'}}"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "response.metadata"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "c281f1f3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'1. Confident: Donald Trump exudes confidence in his speeches, often using phrases like \"I\\'m the only one\" or \"I\\'m the best.\"\\n2. Business-oriented: He frequently highlights his business background, mentioning his companies, properties, and financial successes.\\n3. Opinionated: Trump is not afraid to share his views, often making strong statements about his likes and dislikes.\\n4. Competitive: He presents himself as a winner and emphasizes the importance of winning in various aspects of life.\\n5. Outspoken: Trump is known for his direct and unfiltered communication style, often speaking bluntly and candidly about his thoughts and feelings.'"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "response.response"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "dda2dd0c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The game referred to is the Hunger Games, an annual event in a dystopian society where 24 tributes, young representatives from different districts, fight to the death until only one remains. The contestant described in the context is in the midst of this competition, trying to survive and outsmart her opponents. The Gamemakers, the organizers of the event, manipulate the arena and the tributes' circumstances to create entertainment for the audience. In this part of the story, the contestant faces a wall of fire, precision launchers firing fireballs, and a pack of Careers, the stronger and more ruthless tributes. She must navigate these challenges while avoiding poisonous substances and staying hidden from her enemies.\n"
     ]
    }
   ],
   "source": [
    "response = query_engine.query(\"Summarize the game\")\n",
    "print(response.response)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "a7ea6abd",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'To win the game, it is essential to have a clear understanding of the rules and objectives, maintain a strong focus on your goals, and consistently demonstrate determination and resilience in the face of challenges. Additionally, being adaptable to changing circumstances and learning from past experiences can significantly contribute to your success.'"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "response = query_engine.query(\"How to win the game. Write in 2-3 sentences\")\n",
    "response.response"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5b350f92",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}