examples/notebooks/rag/rag_1D_query_replicate.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Query Data using LLM\n",
    "\n",
    "Here is the overall RAG pipeline.   In this notebook, we will do steps (5), (6), (7), (8), (9)\n",
    "- Importing data is already done in this notebook [rag_1B_load_data_into_milvus.ipynb](rag_1B_load_data_into_milvus.ipynb)\n",
    "- 👉 Step 5: Calculate embedding for user query\n",
    "- 👉 Step 6 & 7: Send the query to vector db to retrieve relevant documents\n",
    "- 👉 Step 8 & 9: Send the query and relevant documents (returned above step) to LLM and get answers to our query\n",
    "\n",
    "![image missing](media/rag-overview-2.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step-1: Configuration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from my_config import MY_CONFIG"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step-2: Load .env file\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "✅ config REPLICATE_API_TOKEN found\n"
     ]
    }
   ],
   "source": [
    "import os,sys\n",
    "## Load Settings from .env file\n",
    "from dotenv import find_dotenv, dotenv_values\n",
    "\n",
    "# _ = load_dotenv(find_dotenv()) # read local .env file\n",
    "config = dotenv_values(find_dotenv())\n",
    "\n",
    "# debug\n",
    "# print (config)\n",
    "\n",
    "MY_CONFIG.REPLICATE_API_TOKEN = config.get('REPLICATE_API_TOKEN')\n",
    "\n",
    "if  MY_CONFIG.REPLICATE_API_TOKEN:\n",
    "    print (\"✅ config REPLICATE_API_TOKEN found\")\n",
    "else:\n",
    "    raise Exception (\"'❌ REPLICATE_API_TOKEN' is not set.  Please set it above to continue...\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step-3: Connect to Vector Database\n",
    "\n",
    "Milvus can be embedded and easy to use.\n",
    "\n",
    "<span style=\"color:blue;\">Note: If you encounter an error about unable to load database, try this: </span>\n",
    "\n",
    "- <span style=\"color:blue;\">In **vscode** : **restart the kernel** of previous notebook. This will release the db.lock </span>\n",
    "- <span style=\"color:blue;\">In **Jupyter**: Do `File --> Close and Shutdown Notebook` of previous notebook. This will release the db.lock</span>\n",
    "- <span style=\"color:blue;\">Re-run this cell again</span>\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "✅ Connected to Milvus instance: ./rag_1_dpk.db\n"
     ]
    }
   ],
   "source": [
    "from pymilvus import MilvusClient\n",
    "\n",
    "milvus_client = MilvusClient(MY_CONFIG.DB_URI)\n",
    "\n",
    "print (\"✅ Connected to Milvus instance:\", MY_CONFIG.DB_URI)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step-4: Setup Embeddings\n",
    "\n",
    "Use the same embeddings we used to index our documents!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/sujee/apps/anaconda3/envs/data-prep-kit-4-021/lib/python3.11/site-packages/sentence_transformers/cross_encoder/CrossEncoder.py:11: TqdmExperimentalWarning: Using `tqdm.autonotebook.tqdm` in notebook mode. Use `tqdm.tqdm` instead to force console mode (e.g. in jupyter console)\n",
      "  from tqdm.autonotebook import tqdm, trange\n",
      "/home/sujee/apps/anaconda3/envs/data-prep-kit-4-021/lib/python3.11/site-packages/huggingface_hub/file_download.py:1142: FutureWarning: `resume_download` is deprecated and will be removed in version 1.0.0. Downloads always resume when possible. If you want to force a new download, use `force_download=True`.\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "from sentence_transformers import SentenceTransformer\n",
    "\n",
    "model = SentenceTransformer(MY_CONFIG.EMBEDDING_MODEL)\n",
    "\n",
    "def get_embeddings (str):\n",
    "    embeddings = model.encode(str, normalize_embeddings=True)\n",
    "    return embeddings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "embeddings len = 384\n",
      "embeddings[:5] =  [ 0.02468893  0.10352131  0.02752644 -0.08551719 -0.01412828]\n"
     ]
    }
   ],
   "source": [
    "# Test embeddings\n",
    "embeddings = get_embeddings('Paris 2024 Olympics')\n",
    "print ('embeddings len =', len(embeddings))\n",
    "print ('embeddings[:5] = ', embeddings[:5])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step-5: Vector Search and RAG"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get relevant documents using vector / sementic search\n",
    "\n",
    "def fetch_relevant_documents (query : str) :\n",
    "    search_res = milvus_client.search(\n",
    "        collection_name=MY_CONFIG.COLLECTION_NAME,\n",
    "        data = [get_embeddings(query)], # Use the `emb_text` function to convert the question to an embedding vector\n",
    "        limit=3,  # Return top 3 results\n",
    "        search_params={\"metric_type\": \"IP\", \"params\": {}},  # Inner product distance\n",
    "        output_fields=[\"text\"],  # Return the text field\n",
    "    )\n",
    "    # print (search_res)\n",
    "\n",
    "    retrieved_docs_with_distances = [\n",
    "        {'text': res[\"entity\"][\"text\"], 'distance' : res[\"distance\"]} for res in search_res[0]\n",
    "    ]\n",
    "    return retrieved_docs_with_distances\n",
    "## --- end ---\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[   {   'distance': 0.5946735143661499,\n",
      "        'text': '3 Model Architecture\\n'\n",
      "                'Table 1: Model configurations for Granite Code models.'},\n",
      "    {   'distance': 0.5919967889785767,\n",
      "        'text': '3 Model Architecture\\n'\n",
      "                'Figure 2: An overview of depth upscaling (Kim et al., 2024) '\n",
      "                'for efficient training of Granite34B-Code. We utilize the 20B '\n",
      "                'model after 1.6T tokens to start training of 34B model with '\n",
      "                'the same code pretraining data without any changes to the '\n",
      "                'training and inference framework.'},\n",
      "    {   'distance': 0.5557882785797119,\n",
      "        'text': 'Granite Code Models: A Family of Open Foundation Models for '\n",
      "                'Code Intelligence\\n'\n",
      "                'Mayank Mishra ⋆ Matt Stallone ⋆ Gaoyuan Zhang ⋆ Yikang Shen '\n",
      "                'Aditya Prasad Adriana Meza Soria Michele Merler Parameswaran '\n",
      "                'Selvam Saptha Surendran Shivdeep Singh Manish Sethi Xuan-Hong '\n",
      "                'Dang Pengyuan Li Kun-Lung Wu Syed Zawad Andrew Coleman '\n",
      "                'Matthew White Mark Lewis Raju Pavuluri Yan Koyfman Boris '\n",
      "                'Lublinsky Maximilien de Bayser Ibrahim Abdelaziz Kinjal Basu '\n",
      "                'Mayank Agarwal Yi Zhou Chris Johnson Aanchal Goyal Hima Patel '\n",
      "                'Yousaf Shah Petros Zerfos Heiko Ludwig Asim Munawar Maxwell '\n",
      "                'Crouse Pavan Kapanipathi Shweta Salaria Bob Calio Sophia Wen '\n",
      "                'Seetharami Seelam Brian Belgodere Carlos Fonseca Amith '\n",
      "                'Singhee Nirmit Desai David D. Cox Ruchir Puri † Rameswar '\n",
      "                'Panda †'}]\n"
     ]
    }
   ],
   "source": [
    "# test relevant vector search\n",
    "import json\n",
    "import pprint\n",
    "\n",
    "question = \"What was the training data used to train Granite models?\"\n",
    "relevant_docs = fetch_relevant_documents(question)\n",
    "pprint.pprint(relevant_docs, indent=4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step-6: Initialize LLM\n",
    "\n",
    "### LLM Choices at Replicate\n",
    "\n",
    "\n",
    "| Model                               | Publisher | Params | Description                                          |\n",
    "|-------------------------------------|-----------|--------|------------------------------------------------------|\n",
    "| ibm-granite/granite-3.0-8b-instruct | IBM       | 8 B    | IBM's newest Granite Model v3.0  (default)           |\n",
    "| ibm-granite/granite-3.0-2b-instruct | IBM       | 2 B    | IBM's newest Granite Model v3.0                      |\n",
    "| meta/meta-llama-3.1-405b-instruct   | Meta      | 405 B  | Meta's flagship 405 billion parameter language model |\n",
    "| meta/meta-llama-3-8b-instruct       | Meta      | 8 B    | Meta's 8 billion parameter language model            |\n",
    "| meta/meta-llama-3-70b-instruct      | Meta      | 70 B   | Meta's 70 billion parameter language model           |\n",
    "\n",
    "References \n",
    "\n",
    "- https://www.ibm.com/granite\n",
    "- https://www.llama.com/\n",
    "- https://replicate.com/  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using model: ibm-granite/granite-3.0-8b-instruct\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "os.environ[\"REPLICATE_API_TOKEN\"] = MY_CONFIG.REPLICATE_API_TOKEN\n",
    "\n",
    "print ('Using model:', MY_CONFIG.LLM_MODEL)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "import replicate\n",
    "\n",
    "def ask_LLM (question, relevant_docs):\n",
    "    context = \"\\n\".join(\n",
    "        [doc['text'] for doc in relevant_docs]\n",
    "    )\n",
    "    print ('============ context (this is the context supplied to LLM) ============')\n",
    "    print (context)\n",
    "    print ('============ end  context ============', flush=True)\n",
    "\n",
    "    system_prompt = \"\"\"\n",
    "    Human: You are an AI assistant. You are able to find answers to the questions from the contextual passage snippets provided.\n",
    "    \"\"\"\n",
    "    user_prompt = f\"\"\"\n",
    "    Use the following pieces of information enclosed in <context> tags to provide an answer to the question enclosed in <question> tags.\n",
    "    <context>\n",
    "    {context}\n",
    "    </context>\n",
    "    <question>\n",
    "    {question}\n",
    "    </question>\n",
    "    \"\"\"\n",
    "\n",
    "    print ('============ here is the answer from LLM... STREAMING... =====')\n",
    "    # The meta/meta-llama-3-8b-instruct model can stream output as it's running.\n",
    "    for event in replicate.stream(\n",
    "        MY_CONFIG.LLM_MODEL,\n",
    "        input={\n",
    "            \"top_k\": 1,\n",
    "            \"top_p\": 0.95,\n",
    "            \"prompt\": user_prompt,\n",
    "            \"max_tokens\": 1024,\n",
    "            \"temperature\": 0.1,\n",
    "            \"system_prompt\": system_prompt,\n",
    "            \"length_penalty\": 1,\n",
    "            # \"max_new_tokens\": 512,\n",
    "            \"stop_sequences\": \"<|end_of_text|>,<|eot_id|>\",\n",
    "            \"prompt_template\": \"<|begin_of_text|><|start_header_id|>system<|end_header_id|>\\n\\n{system_prompt}<|eot_id|><|start_header_id|>user<|end_header_id|>\\n\\n{prompt}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\\n\\n\",\n",
    "            \"presence_penalty\": 0,\n",
    "            \"log_performance_metrics\": False\n",
    "        },\n",
    "    ):\n",
    "        print(str(event), end=\"\")\n",
    "    ## ---\n",
    "    print ('\\n======  end LLM answer ======\\n', flush=True)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step-7: Query"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "============ context (this is the context supplied to LLM) ============\n",
      "3 Model Architecture\n",
      "Table 1: Model configurations for Granite Code models.\n",
      "3 Model Architecture\n",
      "Figure 2: An overview of depth upscaling (Kim et al., 2024) for efficient training of Granite34B-Code. We utilize the 20B model after 1.6T tokens to start training of 34B model with the same code pretraining data without any changes to the training and inference framework.\n",
      "Granite Code Models: A Family of Open Foundation Models for Code Intelligence\n",
      "Mayank Mishra ⋆ Matt Stallone ⋆ Gaoyuan Zhang ⋆ Yikang Shen Aditya Prasad Adriana Meza Soria Michele Merler Parameswaran Selvam Saptha Surendran Shivdeep Singh Manish Sethi Xuan-Hong Dang Pengyuan Li Kun-Lung Wu Syed Zawad Andrew Coleman Matthew White Mark Lewis Raju Pavuluri Yan Koyfman Boris Lublinsky Maximilien de Bayser Ibrahim Abdelaziz Kinjal Basu Mayank Agarwal Yi Zhou Chris Johnson Aanchal Goyal Hima Patel Yousaf Shah Petros Zerfos Heiko Ludwig Asim Munawar Maxwell Crouse Pavan Kapanipathi Shweta Salaria Bob Calio Sophia Wen Seetharami Seelam Brian Belgodere Carlos Fonseca Amith Singhee Nirmit Desai David D. Cox Ruchir Puri † Rameswar Panda †\n",
      "============ end  context ============\n",
      "============ here is the answer from LLM... STREAMING... =====\n",
      "The context does not provide specific details about the training data used to train the Granite models. It only mentions that the 20B model was trained after 1.6T tokens and then used to start training the 34B model with the same code pretraining data. However, it does not specify what this code pretraining data is.\n",
      "======  end LLM answer ======\n",
      "\n",
      "CPU times: user 63.6 ms, sys: 12 ms, total: 75.6 ms\n",
      "Wall time: 1.43 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "question = \"What was the training data used to train Granite models?\"\n",
    "relevant_docs = fetch_relevant_documents(question)\n",
    "ask_LLM(question=question, relevant_docs=relevant_docs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "============ context (this is the context supplied to LLM) ============\n",
      "1 Introduction\n",
      "Attention mechanisms have become an integral part of compelling sequence modeling and transduction models in various tasks, allowing modeling of dependencies without regard to their distance in the input or output sequences [2, 19]. In all but a few cases [27], however, such attention mechanisms are used in conjunction with a recurrent network.\n",
      "3.2 Attention\n",
      "An attention function can be described as mapping a query and a set of key-value pairs to an output, where the query, keys, values, and output are all vectors. The output is computed as a weighted sum\n",
      "7 Conclusion\n",
      "We are excited about the future of attention-based models and plan to apply them to other tasks. We plan to extend the Transformer to problems involving input and output modalities other than text and to investigate local, restricted attention mechanisms to efficiently handle large inputs and outputs such as images, audio and video. Making generation less sequential is another research goals of ours.\n",
      "============ end  context ============\n",
      "============ here is the answer from LLM... STREAMING... =====\n",
      "An attention mechanism is a method used in sequence modeling and transduction models to model dependencies between elements in input or output sequences, regardless of their distance. It maps a query and a set of key-value pairs to an output, which is computed as a weighted sum.\n",
      "======  end LLM answer ======\n",
      "\n",
      "CPU times: user 30.6 ms, sys: 17.3 ms, total: 47.9 ms\n",
      "Wall time: 880 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "question = \"What is attention mechanism?\"\n",
    "relevant_docs = fetch_relevant_documents(question)\n",
    "ask_LLM(question=question, relevant_docs=relevant_docs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "============ context (this is the context supplied to LLM) ============\n",
      "6.1.5 RepoBench, CrossCodeEval: Repository-Level Code Generation\n",
      "StarCoderBase-3B, MBPP = 29.4. StarCoderBase-3B, MBPP+ = 37.8. StableCode-3B, MBPP = 34.8. StableCode-3B, MBPP+ = 43.3. StarCoder2-3B, MBPP = 42.4. StarCoder2-3B, MBPP+ = 48.6. CodeGemma-2B, MBPP = 30.4. CodeGemma-2B, MBPP+ = 30.8. Granite-3B-Code-Base, MBPP = 36.0. Granite-3B-Code-Base, MBPP+ = 45.1. StarCoderBase-7B, MBPP = 34.8. StarCoderBase-7B, MBPP+ = 42.1. CodeLlama-7B, MBPP = 39.0. CodeLlama-7B, MBPP+ = 42.3. StarCoder2-7B, MBPP = 45.4. StarCoder2-7B, MBPP+ = 46.7. CodeGemma-7B, MBPP = 53.0. CodeGemma-7B, MBPP+ = 54.9. Granite-8B-Code-Base, MBPP = 42.2. Granite-8B-Code-Base, MBPP+ = 49.6. StarCoderBase-15B, MBPP = 37.4. StarCoderBase-15B, MBPP+ = 46.1. CodeLlama-13B, MBPP = 30.6. CodeLlama-13B, MBPP+ = 30.1. StarCoder2-15B, MBPP = 51.2. StarCoder2-15B, MBPP+ = 56.6. Granite-20B-Code-Base, MBPP = 43.8. Granite-20B-Code-Base, MBPP+ = 51.6. CodeLlama-34B, MBPP = 48.6. CodeLlama-34B, MBPP+ = 53.6. Granite-34B-Code-Base, MBPP = 47.2. Granite-34B-Code-Base, MBPP+ = 53.1\n",
      "6.1.3 MBPP and MBPP+: Code Generation in Python\n",
      "MBPP (Austin et al., 2021) and MBPP+ (Liu et al., 2023a) are two of the most widely studied benchmarks for evaluating code models. While the prompt for each MBPP problem includes a natural language description followed by a few tests, MBPP+ consists of 35 × more tests than the original benchmarks. We use greedy decoding and report the mean pass@1 for all the models. Table 5 summarizes the results of different base models. As we can see, Granite3B-Code-Base significantly outperforms CodeGemma-2B but falls short of StarCoder2-3B on\n",
      "6.1.4 DS1000: Data Science Tasks in Python\n",
      "The Granite Code models achieve relatively high accuracy across all sizes (e.g., outperforming CodeGemma at 2B-3B scale, StarCoder2 at 7B-8B scale and CodeLlama models with half of the sizes). This shows that our Granite Code models are not only capable of generating good code but also of using libraries more accurately in real data science workflows.\n",
      "============ end  context ============\n",
      "============ here is the answer from LLM... STREAMING... =====\n",
      "I'm sorry, the provided context does not contain information about the moon landing.\n",
      "======  end LLM answer ======\n",
      "\n",
      "CPU times: user 45 ms, sys: 3.19 ms, total: 48.2 ms\n",
      "Wall time: 412 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "question = \"When was the moon landing?\"\n",
    "relevant_docs = fetch_relevant_documents(question)\n",
    "ask_LLM(question=question, relevant_docs=relevant_docs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "data-prep-kit-4-021",
   "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": 4
}