LLM_RAG_few_shot_examples.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "![image](https://raw.githubusercontent.com/IBM/watson-machine-learning-samples/master/cloud/notebooks/headers/watsonx-Prompt_Lab-Notebook.png)\n",
    "# Use Watsonx to respond to natural language questions using RAG approach"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Note:** Please note that for the watsonx challenge, please consider running these notebooks locally on your laptop/desktop."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "This notebook contains the steps and code to demonstrate support of Retrieval Augumented Generation in watsonx.ai. It introduces commands for data retrieval, knowledge base building & querying, and model testing.\n",
    "\n",
    "This notebook uses Python 3.10.\n",
    "\n",
    "#### Objective\n",
    "\n",
    "Use LLM model, Langchain and Milvus to create a Retrieval Augmented Generation (RAG) system. This will allow us to ask questions about our documents (that were not included in the training data), without fine-tunning the Large Language Model (LLM). When using RAG, if you are given a question, you first do a retrieval step to fetch any relevant documents from a special database, a vector database where these documents were indexed.\n",
    "\n",
    "Retrieval Augmented Generation (RAG) is a versatile pattern that can unlock a number of use cases requiring factual recall of information, such as querying a knowledge base in natural language.\n",
    "\n",
    "##### Definitions:\n",
    "\n",
    "* LLM - Large Language Model\n",
    "* Langchain - a framework designed to simplify the creation of applications using LLMs\n",
    "* Vector database - a database that organizes data through high-dimmensional vectors\n",
    "* Milvus - vector database\n",
    "* RAG - Retrieval Augmented Generation (see below more details about RAGs)\n",
    "\n",
    "#### What is a Retrieval Augmented Generation (RAG) system?\n",
    "Large Language Models (LLMs) has proven their ability to understand context and provide accurate answers to various NLP tasks, including summarization, Q&A, when prompted. While being able to provide very good answers to questions about information that they were trained with, they tend to hallucinate when the topic is about information that they do \"not know\", i.e. was not included in their training data. Retrieval Augmented Generation combines external resources with LLMs. The main two components of a RAG are therefore a retriever and a generator.\n",
    "\n",
    "The retriever part can be described as a system that is able to encode our data so that can be easily retrieved the relevant parts of it upon queriying it. The encoding is done using text embeddings, i.e. a model trained to create a vector representation of the information. The best option for implementing a retriever is a vector database. As vector database, there are multiple options, both open source or commercial products. Few examples are ChromaDB, Mevius, FAISS, Pinecone, Weaviate. Our option in this Notebook will be a local instance of ChromaDB (persistent).\n",
    "\n",
    "For the generator part, the obvious option is a LLM. In this Notebook we will use a quantized LLaMA v2 model, from the Kaggle Models collection.\n",
    "\n",
    "The orchestration of the retriever and generator will be done using Langchain. A specialized function from Langchain allows us to create the receiver-generator in one line of code.\n",
    "\n",
    "In its simplest form, RAG requires 3 steps:\n",
    "\n",
    "- Index knowledge base passages (once)\n",
    "- Retrieve relevant passage(s) from knowledge base (for every user query)\n",
    "- Generate a response by feeding retrieved passage into a large language model (for every user query)\n",
    "\n",
    "\n",
    "#### Data:\n",
    "\n",
    "##### Knowledge base:\n",
    "690,000 words worth of cleaned text from Wikipedia.\n",
    "\n",
    "##### Train and Test datasets: \n",
    "There are three question files, one for each year of students: S08, S09, and S10.\n",
    "\n",
    "The \"questionanswerpairs.txt\" files contain both the questions and answers. The columns in this file are as follows:\n",
    "\n",
    "- ArticleTitle is the name of the Wikipedia article from which questions and answers initially came.\n",
    "- Question is the question.\n",
    "- Answer is the answer.\n",
    "- DifficultyFromQuestioner is the prescribed difficulty rating for the question as given to the question-writer.\n",
    "- DifficultyFromAnswerer is a difficulty rating assigned by the individual who evaluated and answered the question, which may differ from the difficulty in field 4.\n",
    "- ArticleFile is the name of the file with the relevant article.\n",
    "\n",
    "Questions that were judged to be poor were discarded from this data set.\n",
    "\n",
    "There are frequently multiple lines with the same question, which appear if those questions were answered by multiple individuals. https://www.kaggle.com/rtatman/questionanswer-dataset\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "<a id=\"setup\"></a>\n",
    "##  Set up the environment\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Install and import dependecies\n",
    "\n",
    "**Note:** For Windows environments, please remove `| tail -n 1` commands in the cell below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "#!pip install sentence_transformers | tail -n 1\n",
    "#!pip install pandas | tail -n 1\n",
    "#!pip install rouge_score | tail -n 1\n",
    "#!pip install nltk | tail -n 1\n",
    "#!pip install \"ibm-watson-machine-learning>=1.0.312\" | tail -n 1\n",
    "#!pip install PyPDF2 | tail -n 1\n",
    "#!pip install langchain | tail -n 1\n",
    "#!pip install --upgrade pip | tail -n 1\n",
    "#!pip install ibm-generative-ai | tail -n 1\n",
    "#!pip install ipywidgets | tail -n 1\n",
    "#!pip install ipywidgets widgetsnbextension pandas-profiling | tail -n 1\n",
    "#!jupyter nbextension enable --py widgetsnbextension\n",
    "#!python3 -m pip install tensorflow-macos | tail -n 1\n",
    "#!python -m pip install tensorflow-metal | tail -n 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# install milvus with [client] extras by pip\n",
    "#!python3 -m pip install \"milvus[client]\" | tail -n 1\n",
    "#!pip install -U --pre torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/nightly/cpu"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Note:** Please restart the notebook kernel to pick up proper version of packages installed above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import pandas as pd\n",
    "import glob\n",
    "import numpy as np\n",
    "import re\n",
    "import string\n",
    "import os\n",
    "import zipfile\n",
    "from pathlib import Path\n",
    "from tqdm.notebook import tqdm\n",
    "import pickle\n",
    "import requests\n",
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "np.random.seed(0)\n",
    "\n",
    "from sentence_transformers import SentenceTransformer\n",
    "from milvus import default_server\n",
    "from pymilvus import (\n",
    "    Collection,\n",
    "    CollectionSchema,\n",
    "    DataType,\n",
    "    FieldSchema,\n",
    "    connections,\n",
    "    utility,\n",
    ")\n",
    "from sklearn.model_selection import train_test_split\n",
    "from langchain.document_loaders import DataFrameLoader\n",
    "from langchain.text_splitter import RecursiveCharacterTextSplitter\n",
    "from langchain.vectorstores import Milvus\n",
    "\n",
    "from genai import Model\n",
    "from genai.model import Credentials\n",
    "from genai.schemas import GenerateParams\n",
    "from dotenv import load_dotenv"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Watsonx API connection\n",
    "This cell defines the credentials required to work with watsonx API for Foundation\n",
    "Model inferencing.\n",
    "\n",
    "**Action:** Provide the IBM Cloud user API key. For details, see\n",
    "[documentation](https://cloud.ibm.com/docs/account?topic=account-userapikey&interface=ui)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "load_dotenv()\n",
    "BAM_API_URL = os.getenv(\"BAM_API_URL\")\n",
    "BAM_API_KEY = os.getenv(\"BAM_API_KEY\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "credentials = Credentials(api_key=BAM_API_KEY, api_endpoint=BAM_API_URL)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Defining the project id\n",
    "The API requires project id that provides the context for the call. We will obtain the id from the project in which this notebook runs. Otherwise, please provide the project id.\n",
    "\n",
    "**Hint**: You can find the `project_id` as follows. Open the prompt lab in watsonx.ai. At the very top of the UI, there will be `Projects / <project name> /`. Click on the `<project name>` link. Then get the `project_id` from Project's Manage tab (Project -> Manage -> General -> Details).\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "try:\n",
    "    project_id = os.getenv('PROJECT_ID')\n",
    "except KeyError:\n",
    "    pproject_id = input(\"Please enter your project_id (hit enter): \")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "<a id=\"data\"></a>\n",
    "## Train and Test data loading"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Load train and test datasets. At first, training dataset (`train_data`) should be used to work with the models to prepare and tune prompt. Then, test dataset (`test_data`) should be used to calculate the metrics score for selected model, defined prompts and parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "---------------\n",
      "(1715, 6)\n",
      "---------------\n",
      "(825, 6)\n",
      "---------------\n",
      "(1458, 6)\n"
     ]
    }
   ],
   "source": [
    "df_08 = pd.read_table('/Users/rrfsantos/Documents/GitHub/portifolio/Generative-AI/S08_question_answer_pairs.txt')\n",
    "df_09 = pd.read_table('/Users/rrfsantos/Documents/GitHub/portifolio/Generative-AI/S09_question_answer_pairs.txt')\n",
    "df_10 = pd.read_table('/Users/rrfsantos/Documents/GitHub/portifolio/Generative-AI/S10_question_answer_pairs.txt', \n",
    "                      encoding = 'windows-1252')  #engine = 'python', error_bad_lines = False)\n",
    "\n",
    "print('-' * 15, df_08.shape, sep='\\n')\n",
    "print('-' * 15, df_09.shape, sep='\\n')\n",
    "print('-' * 15, df_10.shape, sep='\\n')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "---------------\n",
      "(3420, 3)\n"
     ]
    }
   ],
   "source": [
    "df = pd.concat([df_08, df_09, df_10], ignore_index=True)\n",
    "df.drop(['DifficultyFromQuestioner', 'DifficultyFromAnswerer', 'ArticleTitle'], axis = 1, inplace=True)\n",
    "df.dropna(inplace=True)\n",
    "\n",
    "print('-' * 15, df.shape, sep='\\n')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "---------------\n",
      "Question        0\n",
      "Answer          0\n",
      "ArticleFile     0\n",
      "answer_clean    0\n",
      "dtype: int64\n"
     ]
    }
   ],
   "source": [
    "# Limpar coluna \"Answer\"\n",
    "def strip_last_punctuation(s):\n",
    "  if s and s[-1] in string.punctuation:\n",
    "    return s[:-1].strip()\n",
    "  else:\n",
    "    return s.strip()\n",
    "\n",
    "\n",
    "df['answer_clean'] = df['Answer'].str.lower().map(strip_last_punctuation)\n",
    "\n",
    "# Remove os dados faltantes da base de treino\n",
    "df.dropna(inplace=True)\n",
    "\n",
    "print('-' * 15, df.isna().sum(), sep='\\n')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "---------------\n",
      "(2783, 3)\n",
      "---------------\n",
      "Question        0\n",
      "ArticleFile     0\n",
      "answer_clean    0\n",
      "dtype: int64\n"
     ]
    }
   ],
   "source": [
    "df.drop(['Answer'], axis = 1, inplace=True)\n",
    "df.drop_duplicates(subset=['answer_clean', 'Question'], keep='last', inplace = True)\n",
    "print('-' * 15, df.shape, sep='\\n')\n",
    "print('-' * 15, df.isna().sum(), sep='\\n')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_data, test_data = train_test_split(df, test_size=0.3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Build up knowledge base\n",
    "\n",
    "The current state-of-the-art in RAG is to create dense vector representations of the knowledge base in order to calculate the semantic similarity to a given user query.\n",
    "\n",
    "We can generate dense vector representations using embedding models. In this notebook, we use [SentenceTransformers](https://www.google.com/search?client=safari&rls=en&q=sentencetransformers&ie=UTF-8&oe=UTF-8) [all-distilroberta-v1](https://huggingface.co/sentence-transformers/all-distilroberta-v1) to embed both the knowledge base passages and user queries.\n",
    "\n",
    "A vector database is optimized for dense vector indexing and retrieval. This notebook uses [Milvus](https://milvus.io), a user-friendly open-source vector database, licensed under Apache 2.0, which offers good speed and performance with all-MiniLM-L6-v2 embedding model."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "The size of each passage is limited by the embedding model's context window (which is 512 tokens for `all-distilroberta-v1`)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Load knowledge base documents\n",
    "\n",
    "##### Cleaned text from Wikipedia"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Populate a dataframe with information about the docs available \n",
    "\n",
    "files = glob.glob('/Users/rrfsantos/Documents/GitHub/portifolio/Generative-AI/context/*.clean')\n",
    "\n",
    "list_text = []\n",
    "list_file = []\n",
    "\n",
    "for file in files:\n",
    "  with open(file, 'r', encoding = 'utf-8', errors='ignore') as f:\n",
    "    text = f.read()\n",
    "  list_text.append(re.sub(r'\\n+', ' ', text).strip())\n",
    "  list_file.append((os.path.basename(file)).split('.')[0])\n",
    "df = pd.DataFrame({'title': list_file,'text': list_text})\n",
    "\n",
    "missing_cols, missing_rows = ((df.isnull().sum(x) | df.eq(' ').sum(x))\n",
    "        .loc[lambda x: x.gt(0)].index\n",
    "        for x in (0, 1))\n",
    "\n",
    "df.loc[missing_rows, missing_cols]\n",
    "df['id'] = df.reset_index().index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "documents = df.copy()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Create an embedding function\n",
    "\n",
    "Note that you can feed a custom embedding function to be used by Milvus. The performance of Milvus may differ depending on the embedding model used."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Select a Sentence Transformer: https://www.sbert.net/docs/pretrained_models.html\n",
    "embedding_model = SentenceTransformer(\"all-distilroberta-v1\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Start Milvus"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Start the Milvus embedded server"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Server is already running\n"
     ]
    }
   ],
   "source": [
    "if default_server.running != True:\n",
    "    default_server.start()\n",
    "    print(\"Server should have now started\")\n",
    "else:\n",
    "    default_server.stop()\n",
    "    default_server.cleanup()\n",
    "    default_server.start()\n",
    "    print(\"Server is already running\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Establish a connection with the embedded server and print its version information."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "v2.3.3-lite\n"
     ]
    }
   ],
   "source": [
    "connections.connect(host=\"localhost\", port=default_server.listen_port)\n",
    "print(utility.get_server_version())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define the collection"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "COLLECTION_NAME = \"wikipedia_collection\"\n",
    "INDEX_NAME = \"wikipedia_index\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Collection has been deleted\n"
     ]
    }
   ],
   "source": [
    "# Run if you want to drop your old data\n",
    "try:\n",
    "    utility.drop_collection(COLLECTION_NAME)\n",
    "    print(\"Collection has been deleted\")\n",
    "except:  # noqa: E722\n",
    "    pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "id = FieldSchema(\n",
    "    name=\"id\",\n",
    "    dtype=DataType.INT64,\n",
    "    is_primary=True,\n",
    "    auto_id=True,\n",
    ")\n",
    "\n",
    "text = FieldSchema(\n",
    "    name=\"text\",\n",
    "    dtype=DataType.VARCHAR,\n",
    "    max_length=6000,\n",
    ")\n",
    "\n",
    "text_vector = FieldSchema(name=\"text_vector\", dtype=DataType.FLOAT_VECTOR, dim=768)\n",
    "\n",
    "qid = FieldSchema(name=\"qid\", dtype=DataType.INT64)\n",
    "\n",
    "title = FieldSchema(\n",
    "    name=\"title\",\n",
    "    dtype=DataType.VARCHAR,\n",
    "    max_length=5120,\n",
    ")\n",
    "\n",
    "schema = CollectionSchema(\n",
    "    fields=[id, text, text_vector, qid, title],\n",
    "    description=\"SIEM vector store\",\n",
    "    enable_dynamic_field=True,\n",
    ")\n",
    "\n",
    "collection = Collection(\n",
    "    name=COLLECTION_NAME, schema=schema, using=\"default\", shards_num=2\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Embed and index documents with Milvus\n",
    "\n",
    "**Note: Could take several minutes if you don't have pre-built indices**\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Prepare collection"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "text_splitter = RecursiveCharacterTextSplitter(\n",
    "    chunk_size=2500, chunk_overlap=50, length_function=len, add_start_index=False\n",
    ")\n",
    "\n",
    "\n",
    "def split_and_prepare_document_new(qid: str, title: str, text: str):\n",
    "    split_text = text_splitter.split_text(text)\n",
    "    ids = [qid] * len(split_text)\n",
    "    titles = [title] * len(split_text)\n",
    "    embeddings = [\n",
    "        embedding_model.encode(xc, show_progress_bar=False) for xc in split_text\n",
    "    ]\n",
    "    return split_text, ids, titles, embeddings\n",
    "\n",
    "\n",
    "def process_batch(document_list):\n",
    "    batch_results = []\n",
    "\n",
    "    for id, title, text in zip(\n",
    "        document_list[\"id\"].values.tolist(),\n",
    "        document_list[\"title\"].values.tolist(),\n",
    "        document_list[\"text\"].values.tolist(),\n",
    "    ):\n",
    "        for sub_text, sub_id, sub_title, sub_embedding in zip(\n",
    "            *split_and_prepare_document_new(id, title, text)\n",
    "        ):\n",
    "            batch_results.append(tuple((sub_id, sub_title, sub_text, sub_embedding)))\n",
    "    return batch_results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2cb378c8b88b49caae6f656f40011a35",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Processing Documents in Batches:   0%|          | 0/15 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed docs saved to pickle checkpoint\n"
     ]
    }
   ],
   "source": [
    "batch_size = 10\n",
    "processed_docs = []\n",
    "cache_filename = '/Users/rrfsantos/Documents/GitHub/portifolio/prepared-docs.pkl'\n",
    "allow_cache = True\n",
    "\n",
    "if allow_cache and os.path.isfile(cache_filename):\n",
    "    os.remove(cache_filename)\n",
    "\n",
    "for i in tqdm(range(0, len(documents), batch_size), desc=\"Processing Documents in Batches\"):\n",
    "    # find end of batch\n",
    "    i_end = min(i + batch_size, len(documents))\n",
    "    documents_batch = documents[i:i_end]\n",
    "\n",
    "    # Process the batch\n",
    "    processed = process_batch(documents_batch)\n",
    "    processed_docs.extend(processed)\n",
    "\n",
    "with open(cache_filename, \"wb\") as f:\n",
    "    pickle.dump(processed_docs, f)\n",
    "\n",
    "print(\"Processed docs saved to pickle checkpoint\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Insert the embeddings, texts, titles and documents id's in collection."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "10020ff926dc4c64922850a7f4252455",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Inserting documents batches to Milvus VectorDB:   0%|          | 0/2 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "if default_server.running:\n",
    "    collection = Collection(COLLECTION_NAME)\n",
    "    error  = []\n",
    "    batch_size = 1000\n",
    "    for i in tqdm(\n",
    "        range(0, len(processed_docs), batch_size),\n",
    "        desc=\"Inserting documents batches to Milvus VectorDB\",\n",
    "    ):\n",
    "        # find end of batch\n",
    "        i_end = min(i + batch_size, len(processed_docs))\n",
    "        id_l, title_l, text_l, embed_l = list(zip(*processed_docs[i:i_end]))\n",
    "\n",
    "        data_to_insert = [text_l, embed_l, id_l, title_l]\n",
    "        try:\n",
    "            collection.insert(data_to_insert)\n",
    "        except Exception as ex:\n",
    "            print(f\"Failed to insert: {ex}\")\n",
    "            error.append(title_l)\n",
    "else:\n",
    "    print(\"Milvus server is not running! Rerun related notebook cells.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Collection index has been successfully created!\n"
     ]
    }
   ],
   "source": [
    "NLIST_SIZE = 1024\n",
    "\n",
    "index_params = {\n",
    "    \"metric_type\": \"COSINE\",\n",
    "    \"index_type\": \"HNSW\",\n",
    "    \"params\": {\"nlist\": NLIST_SIZE},\n",
    "    \"M\": 16,\n",
    "    \"efConstruction\": 200,\n",
    "}\n",
    "\n",
    "collection.create_index(field_name=\"text_vector\", index_params=index_params)\n",
    "\n",
    "print(\"Collection index has been successfully created!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "<a id=\"models\"></a>\n",
    "## Foundation Models on Watsonx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "load_dotenv()\n",
    "BAM_API_URL = os.getenv(\"BAM_API_URL\")\n",
    "BAM_API_KEY = os.getenv(\"BAM_API_KEY\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "salesforce/codegen2-16b\n",
      "codellama/codellama-34b-instruct\n",
      "tiiuae/falcon-180b\n",
      "tiiuae/falcon-40b\n",
      "ibm/falcon-40b-8lang-instruct\n",
      "google/flan-t5-xl\n",
      "google/flan-t5-xxl\n",
      "google/flan-ul2\n",
      "eleutherai/gpt-neox-20b\n",
      "ibm/granite-13b-chat-v1\n",
      "ibm/granite-13b-chat-v2\n",
      "ibm/granite-13b-instruct-v1\n",
      "ibm/granite-13b-instruct-v2\n",
      "ibm/granite-20b-code-instruct-v1\n",
      "ibm/granite-3b-code-plus-v1\n",
      "elyza/japanese-llama-2-7b-fast\n",
      "elyza/japanese-llama-2-7b-instruct\n",
      "meta-llama/llama-2-13b\n",
      "meta-llama/llama-2-13b-chat\n",
      "meta-llama/llama-2-13b-chat-beam\n",
      "meta-llama/llama-2-70b\n",
      "meta-llama/llama-2-70b-chat\n",
      "thebloke/llama-2-70b-chat-gptq\n",
      "meta-llama/llama-2-7b\n",
      "meta-llama/llama-2-7b-chat\n",
      "mosaicml/mpt-30b\n",
      "ibm/mpt-7b-instruct\n",
      "bigscience/mt0-xxl\n",
      "defog/sqlcoder-34b-alpha\n",
      "bigcode/starcoder\n",
      "flan-t5-xl-mpt-HrlayZEh-2023-10-25-18-15-34\n"
     ]
    }
   ],
   "source": [
    "# get the list of supported models from the API\n",
    "models = Model.models(credentials=credentials)\n",
    "\n",
    "model_ids = []\n",
    "for model_n in models:\n",
    "    print(model_n.id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model was found, it's token limit is 4096.\n"
     ]
    }
   ],
   "source": [
    "# select generative model to use\n",
    "model_id = \"thebloke/llama-2-70b-chat-gptq\"\n",
    "\n",
    "# Iterate over the \"results\" list to find the matching model ID\n",
    "for model_n in models:\n",
    "    if model_n.id == model_id:\n",
    "        model_token_limit = model_n.token_limit\n",
    "        print(f\"Model was found, it's token limit is {model_token_limit}.\")\n",
    "        break\n",
    "else:\n",
    "    print(\"Model was not found, pick a different one!\")\n",
    "    model_token_limit = None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [],
   "source": [
    "# set-up inference parameters\n",
    "params = GenerateParams(decoding_method=\"greedy\",\n",
    "                        max_new_tokens=500,\n",
    "                        min_new_tokens=1,\n",
    "                        repetition_penalty=1,\n",
    "                        temperature=0.5,\n",
    "                        top_k=50,\n",
    "                        top_p=1)\n",
    "\n",
    "model = Model(model=model_id, credentials=credentials, params=params)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The input token limit depends on the selected generative model's max sequence length. The total input tokens in the RAG prompt should not exceed the model's max sequence length minus the number of desired output tokens. The choice of the number of paragraphs to retrieve as context impacts the number tokens in the prompt."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Input token limit: 3595\n"
     ]
    }
   ],
   "source": [
    "# For setting the input token limit we subtract the max_new_tokens (to be generated) and -1 from the model_token_limit\n",
    "input_token_limit = model_token_limit - params.max_new_tokens - 1\n",
    "print(f\"Input token limit: {input_token_limit}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "<a id=\"predict\"></a>\n",
    "## Generate a retrieval-augmented response to a question"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Feed the context and the examples to `watsonx.ai` model."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Feed the context and the examples to `genai` model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Token counting function\n",
    "def token_count(doc):\n",
    "    return model.tokenize([doc])[0].token_count"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`prompt_template` is a function to create a prompt from the given context and question. Changing the prompt will sometimes result in much more appropriate answers (or it may degrade the quality significantly). The prompt template below is most appropriate for short-form extractive use cases.\n",
    "\n",
    "`make_prompt` includes a script to truncate the context length provided as an input in case the total token inputs exceed the model's limit. The paragraphs with the largest distance are truncated first. This functionality is helpful in case the embedded passages are not of the same size."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [],
   "source": [
    "examples = train_data.sample(n=10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [],
   "source": [
    "few_shot_example=[]\n",
    "few_shot_examples=[]\n",
    "\n",
    "for index, row in examples.iterrows():\n",
    "    input = row.Question\n",
    "    output = row.answer_clean\n",
    "    few_shot_example.append(f\"input: {input}\\noutput: {output}\")\n",
    "few_shot_examples='\\n\\n\\n'.join(few_shot_example)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [],
   "source": [
    "def prompt_template(context, few_shot_examples, question_text):\n",
    "    return ('''Below is an instruction that describes a task, paired with an input that provides further context and questions and answer examples.\n",
    "\n",
    "### Input:\n",
    "\n",
    "Based on the context and the examples, write one Response that appropriately answer the Question below.\n",
    "Provide your output in a maximum of 3 sentences.\n",
    "\n",
    "Context:'''\n",
    "            \n",
    "+ \"\\n\\n\"\n",
    "+ f\"{context}\\n\\n\" \n",
    "+ \"##\\n\\n\"\n",
    "+ \"Examples:\\n\\n\"\n",
    "+ f\"{few_shot_examples}\\n\\n\"\n",
    "+ f\"Question: {question_text}\\n\\n\"\n",
    "\n",
    "\"### Response:\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_prompt(question, few_shot_examples, input_token_limit):\n",
    "\n",
    "    ### Create question embedding\n",
    "    question_embeddings = embedding_model.encode(question)\n",
    "\n",
    "    ### Collect the context in vetctordb\n",
    "    \n",
    "    search_params = {\"params\": {\"nprobe\": 10, \"radius\": 0.60}, \"metric_type\": \"COSINE\"} # radius ---> minimum similarity\n",
    "\n",
    "    results = collection.search(\n",
    "       data=[question_embeddings],\n",
    "       anns_field=\"text_vector\",\n",
    "       limit=2,\n",
    "       output_fields=[\"id\",\"text\"],  # name of the field to retrieve from the search result\n",
    "       param=search_params,\n",
    "       round_decimal=-1\n",
    "    )\n",
    "    \n",
    "    print(f\"Found {len(results[0])} results in the collection.\")\n",
    "\n",
    "    documents = []\n",
    "    for raw_result in results:\n",
    "        for result in raw_result:\n",
    "            documents.append(result.entity.get(\"text\"))\n",
    "            #print(\"=========\")\n",
    "            #print(\"Paragraph : \", result.entity.get(\"text\"))\n",
    "            print(\"Distance : \", result.distance)\n",
    "\n",
    "    context = \"\\n\\n\\n\".join(documents)\n",
    "\n",
    "    prompt = prompt_template(context, few_shot_examples, question)\n",
    "\n",
    "    prompt_token_count = token_count(prompt)\n",
    "\n",
    "    if prompt_token_count <= input_token_limit:\n",
    "        return prompt\n",
    "\n",
    "    print(\"exceeded input token limit, truncating context\", prompt_token_count)\n",
    "    \n",
    "    # documents with the lower distance scores are included in the truncated context first\n",
    "    distances = results[0].distances\n",
    "    sorted_indices = sorted(range(len(distances)), key=lambda k: distances[k])\n",
    "\n",
    "    truncated_context = \"\"\n",
    "    token_count_so_far = 0\n",
    "    i = 0\n",
    "\n",
    "    while token_count_so_far <= input_token_limit and i < len(sorted_indices):\n",
    "        doc_index = sorted_indices[i]\n",
    "        document = documents[doc_index]\n",
    "        doc_token_count = token_count(document) + token_count(few_shot_examples)\n",
    "\n",
    "        if token_count_so_far + doc_token_count <= input_token_limit:\n",
    "            truncated_context += document + \"\\n\\n\\n\"\n",
    "            token_count_so_far += doc_token_count\n",
    "        else:\n",
    "            remaining_tokens = input_token_limit - token_count_so_far\n",
    "            truncated_context += document[:remaining_tokens]\n",
    "            break\n",
    "\n",
    "        i += 1\n",
    "\n",
    "    return prompt_template(truncated_context, few_shot_examples, question)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 2 µs, sys: 1e+03 ns, total: 3 µs\n",
      "Wall time: 3.81 µs\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Found 1 results in the collection.\n",
      "Distance :  0.6147989630699158\n",
      "Found 0 results in the collection.\n",
      "Found 0 results in the collection.\n",
      "Found 0 results in the collection.\n",
      "Found 2 results in the collection.\n",
      "Distance :  0.6273016929626465\n",
      "Distance :  0.6147995591163635\n",
      "Found 2 results in the collection.\n",
      "Distance :  0.6437525153160095\n",
      "Distance :  0.6328293681144714\n",
      "Found 2 results in the collection.\n",
      "Distance :  0.666768491268158\n",
      "Distance :  0.6583305597305298\n",
      "Found 0 results in the collection.\n",
      "Found 0 results in the collection.\n",
      "Found 2 results in the collection.\n",
      "Distance :  0.6847279071807861\n",
      "Distance :  0.6610898971557617\n",
      "Found 1 results in the collection.\n",
      "Distance :  0.6345176100730896\n",
      "Found 1 results in the collection.\n",
      "Distance :  0.6066610813140869\n",
      "Found 0 results in the collection.\n",
      "Found 0 results in the collection.\n",
      "Found 0 results in the collection.\n",
      "Found 0 results in the collection.\n",
      "Found 1 results in the collection.\n",
      "Distance :  0.601598858833313\n",
      "Found 2 results in the collection.\n",
      "Distance :  0.6914989948272705\n",
      "Distance :  0.6598888635635376\n",
      "Found 0 results in the collection.\n",
      "Found 2 results in the collection.\n",
      "Distance :  0.6179383397102356\n",
      "Distance :  0.6122869253158569\n"
     ]
    }
   ],
   "source": [
    "%time\n",
    "questions = test_data.sample(n=20)\n",
    "questions = questions['Question'].tolist()\n",
    "collection = Collection(COLLECTION_NAME)      # Get the existing collection.\n",
    "collection.load() \n",
    "prompt_texts = []\n",
    "for question in questions:\n",
    "    prompt_text = make_prompt(question, few_shot_examples, input_token_limit)\n",
    "    prompt_texts.append(prompt_text)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Generate responses"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 2 µs, sys: 1e+03 ns, total: 3 µs\n",
      "Wall time: 9.06 µs\n"
     ]
    }
   ],
   "source": [
    "%time\n",
    "answers = []\n",
    "for response in model.generate(prompt_texts):\n",
    "    answers.append(response.generated_text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Below is an instruction that describes a task, paired with an input that provides further context and questions and answer examples.\n",
      "\n",
      "### Input:\n",
      "\n",
      "Based on the context and the examples, write one Response that appropriately answer the Question below.\n",
      "Provide your output in a maximum of 3 sentences.\n",
      "\n",
      "Context:\n",
      "\n",
      "elephant Elephants (Elephantidae) are a family in the order Proboscidea in the class Mammalia. They were once classified along with other thick skinned animals in a now invalid order, Pachydermata. There are three living species: the African Bush Elephant, the African Forest Elephant (until recently known collectively as the African Elephant), and the Asian Elephant (also known as the Indian Elephant). Other species have become extinct since the last ice age, which ended about 10,000 years ago, the Mammoth being the most well-known of these. The word \"elephant\" has its origins in the Greek ἐλέφας, meaning \"ivory\" or \"elephant\". Elephants are mammals, and the largest land animals alive today.    The elephant's gestation period is 22 months, the longest of any land animal. At birth it is common for an elephant calf to weigh 120 kilograms (265 lb).  An elephant may live as long as 70 years, sometimes longer. The largest elephant ever recorded was shot in Angola in 1956. This male weighed about 12,000 kg (26,400 lb),    with a shoulder height of 4.2 m (13.8 ft), a metre (3 ft 4 in) taller than the average male African elephant.    The smallest elephants, about the size of a calf or a large pig, were a prehistoric species that lived on the island of Crete during the Pleistocene epoch. Bate, D.M.A. 1907. On Elephant Remains from Crete, with Description of Elephas creticus sp.n. Proc. zool. Soc. London: 238-250. Elephants are symbols of wisdom in Asian cultures, and are famed for their memory and high intelligence, and are thought to be on par with cetaceans    and hominids   . Aristotle once said the elephant was \"the beast which passeth all others in wit and mind.\" Elephants are increasingly threatened by human intrusion and poaching. Once numbering in the millions, the African elephant population has dwindled to between 470,000 and 690,000 individuals.    The elephant is now a protected species worldwide, with restrictions in place on capture, domestic use, and trade in products such as ivory. Elephants generally have no natural predators, although lions may take calves and occasionally adults. A. J. Loveridge, J. E. Hunt, F. Murindagomo   D. W. Macdonald. (2006) Influence of drought on predation of elephant (Loxodonta africana) calves by lions (Panthera leo) in an African wooded savannah. Journal of Zoology 270:3, 523–530     In some areas, lions may regularly take to preying on elephants. Hemson, Graham (2003) The Ecology and Conservation of Lions:\n",
      "\n",
      "##\n",
      "\n",
      "Examples:\n",
      "\n",
      "input: What is an economic use of a duck?\n",
      "output: being farmed for their meat\n",
      "\n",
      "\n",
      "input: What is cause of polar bear's skin diseases?\n",
      "output: mites or other parasites\n",
      "\n",
      "\n",
      "input: Why did Bequerel win the Nobel Prize in Physics?\n",
      "output: discovery of spontaneous radioactivity\n",
      "\n",
      "\n",
      "input: Do penguins live almost exclusively in the Southern Hemisphere?\n",
      "output: yes\n",
      "\n",
      "\n",
      "input: How heavy was the largest polar bear on record?\n",
      "output: he allegedly weighed 1002 kg\n",
      "\n",
      "\n",
      "input: Is the leopard an agile and graceful predator?\n",
      "output: yes\n",
      "\n",
      "\n",
      "input: Can citizens propose changes to the Constitution?\n",
      "output: yes\n",
      "\n",
      "\n",
      "input: Is Korean the official language of Korea?\n",
      "output: yes\n",
      "\n",
      "\n",
      "input: Do most Japanese people employ politeness?\n",
      "output: yes\n",
      "\n",
      "\n",
      "input: Where was Celsius born?\n",
      "output: uppsala in sweden\n",
      "\n",
      "Question: Where was the largest elephant ever recorded shot?\n",
      "\n",
      "### Response:\n"
     ]
    }
   ],
   "source": [
    "print(prompt_texts[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_result = pd.DataFrame({'Question': questions, 'Answer': answers})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Where was the largest elephant ever recorded shot?\n",
      "\n",
      "\n",
      "The largest elephant ever recorded was shot in Angola in 1956.\n"
     ]
    }
   ],
   "source": [
    "print(df_result.Question[0])\n",
    "print(df_result.Answer[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "metadata": {},
   "outputs": [],
   "source": [
    "#default_server.stop()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}