notebooks/qa-gen-query-langchain.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# LangChain and Llama2: Prompt Templates and Batch GPU Inference\n",
    "\n",
    "This notebook uses LangChain and local Llama2-Chat inference that can be run on consumer grade hardware. The following LangChain features explored are:\n",
    "1) [LangChain Custom Prompt Template](https://python.langchain.com/docs/modules/model_io/prompts/prompt_templates/custom_prompt_template) for a Llama2-Chat model\n",
    "2) [Hugging Face Local Pipelines](https://python.langchain.com/docs/integrations/llms/huggingface_pipelines)\n",
    "3) [4-Bit Quantization](https://huggingface.co/blog/4bit-transformers-bitsandbytes)\n",
    "4) [Batch GPU Inference](https://python.langchain.com/docs/integrations/llms/huggingface_pipelines#batch-gpu-inference)\n",
    "\n",
    "Llama2-chat was trained using the below [template](https://gpus.llm-utils.org/llama-2-prompt-template/) and should be prompted the same to best performance. \n",
    "\n",
    "**NOTE**: `<s>` is the beginning of sequence (bos) token.\n",
    "\n",
    "**Llama2-Chat Prompt Template**\n",
    "```\n",
    "<s>[INST] <<SYS>>\n",
    "{your_system_message}\n",
    "<</SYS>>\n",
    "\n",
    "{user_message_1} [/INST]\n",
    "```\n",
    "\n",
    "At the time of writing LangChain does not offer a Llama2-Chat prompt template but Custom Prompt Templates be created for such situations."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# User Inputs and Libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Example context\n",
    "EXAMPLE_CONTEXT = \"\"\" \n",
    "The Eiffel Tower is a wrought-iron lattice tower on the Champ de Mars in Paris, France. \n",
    "It is named after the engineer Gustave Eiffel, whose company designed and built the tower from 1887 to 1889.\n",
    "The Eiffel Tower is 1,083 ft tall.\n",
    "\"\"\".strip()\n",
    "\n",
    "# Llama2-Chat Prompt Template\n",
    "llama2_template = \"\"\"\n",
    "<s>[INST] <<SYS>>\n",
    "{your_system_message}\n",
    "<</SYS>>\n",
    "\n",
    "{user_message_1} [/INST]\n",
    "\"\"\".strip()\n",
    "\n",
    "# System Message\n",
    "sys_template = \"\"\"\n",
    "You are a question generating assistant. \n",
    "Given a document, please generate a simple and short question based on the information provided.\n",
    "The question can be a maximum of 10 words long.\n",
    "\"\"\".strip()\n",
    "\n",
    "# Human Message\n",
    "human_template = \"\"\"\n",
    "DOCUMENT: {context}\n",
    "QUESTION: Your question here.\n",
    "\"\"\".strip()\n",
    "\n",
    "# Path to Model\n",
    "model_id = '/nvme4tb/Projects/llama2_models/Llama-2-13b-chat-hf'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import libraries and packages\n",
    "import gc\n",
    "import torch\n",
    "from time import time\n",
    "from torch import cuda, bfloat16\n",
    "from transformers import (AutoConfig,\n",
    "                          AutoTokenizer,\n",
    "                          AutoModelForCausalLM,\n",
    "                          BitsAndBytesConfig,\n",
    "                          pipeline)\n",
    "from langchain.llms import HuggingFacePipeline\n",
    "from langchain.prompts import (ChatPromptTemplate,\n",
    "                               HumanMessagePromptTemplate,\n",
    "                               SystemMessagePromptTemplate,\n",
    "                               StringPromptTemplate)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# LangChain Standard Chat Prompt Templates\n",
    "\n",
    "LangChain [SystemMessagePromptTemplate](https://api.python.langchain.com/en/latest/prompts/langchain.prompts.chat.SystemMessagePromptTemplate.html#langchain.prompts.chat.SystemMessagePromptTemplate) and [HumanMessagePromptTemplate](https://api.python.langchain.com/en/latest/prompts/langchain.prompts.chat.HumanMessagePromptTemplate.html#langchain.prompts.chat.HumanMessagePromptTemplate) are commonly used LangChain prompt templates; however, at the time of writing, are not optimized for Llama2-Chat.\n",
    "\n",
    "The below cell shows the prompt template returned using these LangChain classes.\n",
    "\n",
    "**REFERENCES**\n",
    "- [LangChain Message Prompt Templates](https://python.langchain.com/docs/modules/model_io/prompts/prompt_templates/msg_prompt_templates)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "System: You are a question generating assistant. \n",
      "Given a document, please generate a simple and short question based on the information provided.\n",
      "The question can be a maximum of 10 words long.\n",
      "Human: DOCUMENT: The Eiffel Tower is a wrought-iron lattice tower on the Champ de Mars in Paris, France. \n",
      "It is named after the engineer Gustave Eiffel, whose company designed and built the tower from 1887 to 1889.\n",
      "The Eiffel Tower is 1,083 ft tall.\n",
      "QUESTION: Your question here.\n"
     ]
    }
   ],
   "source": [
    "# Chat template\n",
    "chat_template = ChatPromptTemplate.from_messages(\n",
    "    [SystemMessagePromptTemplate.from_template(sys_template),\n",
    "     HumanMessagePromptTemplate.from_template(human_template)])\n",
    "\n",
    "# Invoke the chat template\n",
    "chat_invoked = chat_template.invoke({'context': EXAMPLE_CONTEXT})\n",
    "\n",
    "# Print the template that would be passed to the LLM\n",
    "print(chat_invoked.to_string())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notice the above prompt template does not match what is required for Llama2-chat. In the next section will implement a LangChain custom prompt template that can be used for Llama2-Chat."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Custom Prompt Template\n",
    "\n",
    "This section demonstrates how to create a [LangChain custom prompt template](https://python.langchain.com/docs/modules/model_io/prompts/prompt_templates/custom_prompt_template) for Llama2. The custom class could easily be modified to work with any LLM model of choice. Further, the input parameters (e.g., model_template, system_message, human_message) could be pointed to a prompt template databases for robust usage. \n",
    "\n",
    "**References**\n",
    "- [LangChain - Custom Agent With Tool Retrieval](https://python.langchain.com/docs/modules/agents/how_to/custom_agent_with_tool_retrieval)\n",
    "- [LangChain - Custom Prompt Template](https://python.langchain.com/docs/modules/model_io/prompts/prompt_templates/custom_prompt_template)\n",
    "- [LangChain - StringPromptTemplate](https://api.python.langchain.com/en/latest/prompts/langchain.prompts.base.StringPromptTemplate.html#langchain.prompts.base.StringPromptTemplate)\n",
    "- [Blog on LangChain and Llama2](https://www.mlexpert.io/prompt-engineering/langchain-quickstart-with-llama-2)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<s>[INST] <<SYS>>\n",
      "You are a question generating assistant. \n",
      "Given a document, please generate a simple and short question based on the information provided.\n",
      "The question can be a maximum of 10 words long.\n",
      "<</SYS>>\n",
      "\n",
      "DOCUMENT: The Eiffel Tower is a wrought-iron lattice tower on the Champ de Mars in Paris, France. \n",
      "It is named after the engineer Gustave Eiffel, whose company designed and built the tower from 1887 to 1889.\n",
      "The Eiffel Tower is 1,083 ft tall.\n",
      "QUESTION: Your question here. [/INST]\n"
     ]
    }
   ],
   "source": [
    "# Configure a class that defines a function to explain the source code of the given function\n",
    "class Llama2ChatPromptTemplate(StringPromptTemplate):\n",
    "    \"\"\"\n",
    "    Llama2-Chat prompt template customized for system and human messages\n",
    "    \"\"\"\n",
    "    # Define templates\n",
    "    model_template: str\n",
    "    system_template: str\n",
    "    user_template: str\n",
    "\n",
    "\n",
    "    def __get_template(self,\n",
    "                       model_template: str,\n",
    "                       your_system_message: str,\n",
    "                       user_message_1: str) -> str:\n",
    "        \"\"\"\n",
    "        Insert the System and User Messages into the Model Prompt Template\n",
    "\n",
    "        Args:\n",
    "            model_template (str): Model prompt template (e.g. Llama2-Chat)\n",
    "            your_system_message (str): System Message with placeholders for examples, etc.\n",
    "            user_message_1 (str): User message with placeholders for context, etc.\n",
    "\n",
    "        Returns:\n",
    "            str: Prompt template with placeholders (context, documents, examples, etc.)\n",
    "        \"\"\"\n",
    "        # Insert system message into model template, then insert human message\n",
    "        template = model_template.replace('{your_system_message}', your_system_message)\n",
    "        template = template.replace('{user_message_1}', user_message_1)\n",
    "        return template\n",
    "\n",
    "\n",
    "    def format(self, **kwargs) -> str:\n",
    "        \"\"\"\n",
    "        LangChain method for formatting the template\n",
    "        \"\"\"\n",
    "        # Create a prompt template with placeholder for context\n",
    "        PROMPT = self.__get_template(model_template=self.model_template,\n",
    "                                     your_system_message=self.system_template,\n",
    "                                     user_message_1=self.user_template)\n",
    "        \n",
    "        # Generate the prompt to be sent to the llm\n",
    "        prompt = PROMPT.format(context=kwargs['context'])\n",
    "        return prompt\n",
    "\n",
    "# Initialize Prompt Template\n",
    "llama2_prompt_template = Llama2ChatPromptTemplate(model_template=llama2_template,\n",
    "                                                  system_template=sys_template,\n",
    "                                                  user_template=human_template,\n",
    "                                                  input_variables=[\"context\"])\n",
    "\n",
    "# Create the prompt using example context\n",
    "prompt = llama2_prompt_template.format(context=EXAMPLE_CONTEXT)\n",
    "\n",
    "# Print the template that would be passed to the LLM\n",
    "print(prompt)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The above template satisfies the Llama2-Chat format. Again, the above class can be modified to suit any LLM."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Chat with Llama2\n",
    "\n",
    "The following will form a LangChain [Chain](https://python.langchain.com/docs/modules/chains/) using their [LangChain Expression Language (LCEL)](https://python.langchain.com/docs/expression_language/). First, a `text-generation` model will be configured using 4-bit quantization and then the the above prompt template will be chained to the model. \n",
    "\n",
    "This will provide a chat model that can be invoked and queries will be generated.\n",
    "\n",
    "**REFERENCES**\n",
    "- [LangChain HuggingFace Pipeline](https://python.langchain.com/docs/integrations/llms/huggingface_pipelines)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "09d8a88666bb4324803495e9178aa108",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Select the device\n",
    "device = f'cuda:{cuda.current_device()}' if cuda.is_available() else 'cpu'\n",
    "\n",
    "# Set quantization configuration to load large model with less GPU memory\n",
    "# this requires the `bitsandbytes` library\n",
    "bnb_config = BitsAndBytesConfig(\n",
    "    load_in_4bit=True,\n",
    "    bnb_4bit_quant_type='nf4',\n",
    "    bnb_4bit_use_double_quant=True,\n",
    "    bnb_4bit_compute_dtype=bfloat16\n",
    ")\n",
    "\n",
    "# Model\n",
    "model_config = AutoConfig.from_pretrained(model_id)\n",
    "model = AutoModelForCausalLM.from_pretrained(\n",
    "    model_id,\n",
    "    trust_remote_code=True,\n",
    "    config=model_config,\n",
    "    quantization_config=bnb_config,\n",
    "    device_map='auto',\n",
    ")\n",
    "model.eval() # set to evaluation for inference only\n",
    "\n",
    "# Tokenizer\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_id)\n",
    "\n",
    "# Transformers pipeline\n",
    "pipe = pipeline(\n",
    "    model=model, \n",
    "    tokenizer=tokenizer,\n",
    "    return_full_text=True,  # langchain expects the full text\n",
    "    task='text-generation',\n",
    "    # we pass model parameters here too\n",
    "    # stopping_criteria=stopping_criteria,  # without this model rambles during chat\n",
    "    temperature=0.1,  # 'randomness' of outputs, 0.0 is the min and 1.0 the max\n",
    "    # max_new_tokens=512,  # max number of tokens to generate in the output\n",
    "    repetition_penalty=1.1  # without this output begins repeating\n",
    ")\n",
    "\n",
    "# LangChain Hugging Face Pipeline\n",
    "hf = HuggingFacePipeline(pipeline=pipe)\n",
    "\n",
    "# Create a chain\n",
    "chain = llama2_prompt_template | hf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Sun Nov 26 13:40:13 2023       \n",
      "+-----------------------------------------------------------------------------+\n",
      "| NVIDIA-SMI 520.61.05    Driver Version: 520.61.05    CUDA Version: 11.8     |\n",
      "|-------------------------------+----------------------+----------------------+\n",
      "| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |\n",
      "| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |\n",
      "|                               |                      |               MIG M. |\n",
      "|===============================+======================+======================|\n",
      "|   0  NVIDIA GeForce ...  On   | 00000000:01:00.0 Off |                  N/A |\n",
      "| 33%   62C    P2   116W / 350W |   7383MiB / 24576MiB |      0%      Default |\n",
      "|                               |                      |                  N/A |\n",
      "+-------------------------------+----------------------+----------------------+\n",
      "|   1  NVIDIA GeForce ...  On   | 00000000:02:00.0 Off |                  N/A |\n",
      "| 30%   60C    P2   120W / 350W |   9031MiB / 24576MiB |      0%      Default |\n",
      "|                               |                      |                  N/A |\n",
      "+-------------------------------+----------------------+----------------------+\n",
      "                                                                               \n",
      "+-----------------------------------------------------------------------------+\n",
      "| Processes:                                                                  |\n",
      "|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |\n",
      "|        ID   ID                                                   Usage      |\n",
      "|=============================================================================|\n",
      "|    0   N/A  N/A    769801      C   ...GenQuery/.venv/bin/python     3368MiB |\n",
      "|    0   N/A  N/A   1000141      C   ...GenQuery/.venv/bin/python     4012MiB |\n",
      "|    1   N/A  N/A    769801      C   ...GenQuery/.venv/bin/python     4042MiB |\n",
      "|    1   N/A  N/A   1000141      C   ...GenQuery/.venv/bin/python     4986MiB |\n",
      "+-----------------------------------------------------------------------------+\n"
     ]
    }
   ],
   "source": [
    "# View GPU Memory\n",
    "!nvidia-smi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  Sure! Here's a simple and short question based on the information provided:\n",
      "\n",
      "\"What is the height of the Eiffel Tower in Paris?\"\n"
     ]
    }
   ],
   "source": [
    "# Invoke the chain and get a response from Llama2\n",
    "print(chain.invoke({\"context\": EXAMPLE_CONTEXT}))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As shown above a question is returned based on the `EXAMPLE_CONTEXT` provided to Llama2."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Generate Context-Query Pairs using Batch GPU Inference\n",
    "\n",
    "The HuggingFacePipeline can be used to Batch GPU Inferences and it will be demonstrated on a subset of the [Stanford Question Answering Dataset squad_v2](https://huggingface.co/datasets/squad_v2)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load SquadV2 Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "df.shape: (100, 3)\n",
      "Columns: ['id', 'context', 'question']\n"
     ]
    }
   ],
   "source": [
    "from datasets import load_dataset\n",
    "from pathlib import Path\n",
    "from types import SimpleNamespace\n",
    "\n",
    "# Specify paths to data, prompt templates, llama model, etc.\n",
    "paths = {'base_dir': Path.cwd().parents[0],\n",
    "         'prompt_vanilla': 'notebooks/question-answering-prompts/vanilla.txt',\n",
    "         'prompt_gbq': 'notebooks/question-answering-prompts/gbq.txt',\n",
    "         'squad_data': 'data/squad_v2',\n",
    "         'model': '/nvme4tb/Projects/llama2_models/Llama-2-13b-chat-hf',\n",
    "         }\n",
    "\n",
    "# Number of context samples for experimentation\n",
    "NUM_SAMPLES = 100\n",
    "\n",
    "# Convert from dictionary to SimpleNamespace\n",
    "paths = SimpleNamespace(**paths)\n",
    "\n",
    "# Load squad_v2 data locally from disk\n",
    "df = load_dataset(str(paths.base_dir / paths.squad_data),\n",
    "                  split='train').to_pandas()\n",
    "\n",
    "# Remove redundant context\n",
    "df = df.drop_duplicates(subset=['context', 'title']).reset_index(drop=True)\n",
    "\n",
    "# Randomly select 50 contexts\n",
    "df = df.sample(n=NUM_SAMPLES, random_state=42)[['id', 'context', 'question']]\n",
    "\n",
    "# Print Info.\n",
    "print(f'df.shape: {df.shape}')\n",
    "print(f'Columns: {df.columns.tolist()}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Update the hf pipeline batch size\n",
    "hf.batch_size = 50\n",
    "\n",
    "# Create a chain\n",
    "chain_batch = llama2_prompt_template | hf.bind(stop=['\\n\\n'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Place the SquadV2 context in a list of dictionaries\n",
    "contexts = []\n",
    "for context in df.context.tolist():\n",
    "    contexts.append({'context': context})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total Time to Generate 100 Queries: 4.5 mins.\n",
      "Avg. Amount of Seconds Per Sample: 2.7\n"
     ]
    }
   ],
   "source": [
    "# Start time\n",
    "st = time()\n",
    "\n",
    "# GPU Batch Inference\n",
    "queries = chain_batch.batch(contexts)\n",
    "\n",
    "# Total time to generate the queries\n",
    "total_secs = time() - st\n",
    "secs_per_sample = (total_secs / NUM_SAMPLES)\n",
    "print(f'Total Time to Generate {NUM_SAMPLES} Queries: {(total_secs / 60):.1f} mins.')\n",
    "print(f'Avg. Amount of Seconds Per Sample: {secs_per_sample:.1f}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Example # 1\n",
      "Context: The Oklahoma City Police Department, has a uniformed force of 1,169 officers and 300+ civilian employees. The Department has a central police station and five substations covering 2,500 police reporting districts that average 1/4 square mile in size.\n",
      "Synthetic Query: How many substations does Oklahoma city have?\n",
      "\n",
      "Example # 2\n",
      "Context: The U.S. Federal Reserve and central banks around the world have taken steps to expand money supplies to avoid the risk of a deflationary spiral, in which lower wages and higher unemployment lead to a self-reinforcing decline in global consumption. In addition, governments have enacted large fiscal stimulus packages, by borrowing and spending to offset the reduction in private sector demand caused by the crisis. The U.S. Federal Reserve's new and expanded liquidity facilities were intended to enable the central bank to fulfill its traditional lender-of-last-resort role during the crisis while mitigating stigma, broadening the set of institutions with access to liquidity, and increasing the flexibility with which institutions could tap such liquidity.\n",
      "Synthetic Query: What have central banks around the world done to avoid the risk of a deflationary spiral?\n",
      "\n",
      "Example # 3\n",
      "Context: The two finalists were Kris Allen and Adam Lambert, both of whom had previously landed in the bottom three at the top five. Allen won the contest in the most controversial voting result since season two. It was claimed, later retracted, that 38 million of the 100 million votes cast on the night came from Allen's home state of Arkansas alone, and that AT&T employees unfairly influenced the votes by giving lessons on power-texting at viewing parties in Arkansas.\n",
      "Synthetic Query: Who were the final two contestants on season eight of American Idol?\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Assign to dataframe and strip any whitespaces\n",
    "df['query'] = [x.strip() for x in queries]\n",
    "\n",
    "# View a few examples\n",
    "for ii in range(3):\n",
    "    print(f'Example # {ii + 1}')\n",
    "    print(f'Context: {df.iloc[ii].context}')\n",
    "    print(f'Synthetic Query: {df.iloc[ii].question}\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook showcased several LangChain features and is generic so other models or prompt templates can be easily integrated for your application."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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.10.12"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}