not_so_simple_phi2_chatbot.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "2f4717d2-352e-461e-95f9-24e433ef953b",
   "metadata": {},
   "source": [
    "# Not-so-simple phi2 chatbot\n",
    "\n",
    "This notebook uses all the pieces explained in\n",
    "\n",
    "+ `experiments.ipyn` and\n",
    "+ `experiments_langchain.ipyn`\n",
    "\n",
    "to implement a chatbot that can\n",
    "\n",
    "+ classify input queries into 3 different categories (`support`, `sales` and `joke`),\n",
    "+ respond accordingly to each of the user queries depending on the category they belong to."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a4806ab4-bcb1-4cb0-93bd-886b94651227",
   "metadata": {},
   "source": [
    "## Putting it all together\n",
    "\n",
    "Now, in order to put it all together (prompting, conditional branching and conversation memory), we first need to modify a bit our prompts from the _routing_ section. They were meant to be used standalone and they make use of the `Instruction/Output` template from `phi-2`. When we add the conversation bits, this doesn't make sense anymore.\n",
    "\n",
    "We're going to do the following steps:\n",
    "1. Load the model\n",
    "2. Create the classification chain\n",
    "3. Add the branch template\n",
    "4. Add the conversation memory"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6a73ffcd-2eef-47bf-9581-121fc50d3106",
   "metadata": {},
   "source": [
    "### Step 1 load the model\n",
    "\n",
    "The very first thing we have to do is to load the phi model in a format that can be used with LangChain. We've created some conveniency functions for that in `load_phi_model.py`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "b6f4136d-92af-4a37-ba33-60c491938e95",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/grego/code/jobsearch/ipglobal-2/env/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Your device is cuda\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading checkpoint shards: 100%|██████████████████████████████████████████████| 2/2 [00:00<00:00,  2.67it/s]\n",
      "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n"
     ]
    }
   ],
   "source": [
    "from load_phi_model import load_phi_model_and_tokenizer, get_langchain_model\n",
    "\n",
    "model, tokenizer = load_phi_model_and_tokenizer()\n",
    "hf = get_langchain_model(model, tokenizer)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "675cad63-11c2-4f59-a849-d74706b36d12",
   "metadata": {},
   "source": [
    "### Step 2: create the initial classification template\n",
    "\n",
    "We change `text` variable to `human_input` as this is the name we'll use for the chat at the end."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "95fccc17-3f83-4983-bced-b1e327a00b52",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.prompts import PromptTemplate, ChatPromptTemplate\n",
    "from langchain_core.output_parsers import StrOutputParser\n",
    "from langchain_core.runnables import RunnableBranch, RunnableLambda, RunnablePassthrough\n",
    "from operator import itemgetter\n",
    "from langchain.memory import ConversationBufferMemory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "b2e38fda-abdb-4739-b488-857c2172ee95",
   "metadata": {},
   "outputs": [],
   "source": [
    "classification_template = \"\"\"\n",
    "Instruct: Classify the following text in one of the following categories: [\"support\", \"sales\", \"joke\"]. Output only the name of the category.\n",
    "+ \"support\" for customer support texts\n",
    "+ \"sales\" for sales and comercial texts\n",
    "+ \"joke\" for jokes, funny or comedy like texts\n",
    "Text: {human_input}\n",
    "Output:\n",
    "\"\"\".strip()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "48b05df2-a837-4556-815e-93d3f25bd121",
   "metadata": {},
   "outputs": [],
   "source": [
    "classification_prompt = ChatPromptTemplate.from_template(classification_template)\n",
    "classification_chain = (\n",
    "    classification_prompt\n",
    "    | hf\n",
    "    | StrOutputParser()\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "59fb28d2-20a5-4579-abaf-9c06451ab4f0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " support\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(classification_chain.invoke({\"human_input\": \"Can I track my order? I'm eager to know its status\"}))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "af7248da-e408-4591-9dce-1cc9a109f6b5",
   "metadata": {},
   "source": [
    "### Step 3: add the branch template\n",
    "\n",
    "The branch template gives the right instructions for the chatbot. Depending on the result of the classification, we provide different instructions to the language model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "016f53e3-3b1a-4232-ae04-bc4f87dee331",
   "metadata": {},
   "outputs": [],
   "source": [
    "support_instructions = \"\"\"\\\n",
    "You are a customer support agent. It seems that the user may have some issues. Answer to their query politely and sincerely. \\\n",
    "Be kind, understanding and say you're sorry for the inconvenience or the situation whenever necessary. Be brief and to the point.\\\n",
    "\"\"\"\n",
    "\n",
    "sales_instructions = \"\"\"\\\n",
    "You are an aggressive salesperson. The user is looking for some information on products. \\\n",
    "Reply to their query by giving information on related products and showcasing how good they are and why they should buy them. \\\n",
    "Be brief and to the point.\n",
    "\"\"\"\n",
    "\n",
    "joke_instructions = \"\"\"\\\n",
    "You are a comedian. The user want's to have some fun. Reply to their query in a funny way.\\\n",
    "\"\"\"\n",
    "\n",
    "general_instructions = \"\"\"\\\n",
    "Instruction: Respond to the following query.\\\n",
    "\"\"\"\n",
    "\n",
    "support_chain = PromptTemplate.from_template(support_instructions)\n",
    "sales_chain = PromptTemplate.from_template(sales_instructions)\n",
    "joke_chain = PromptTemplate.from_template(joke_instructions)\n",
    "general_chain = PromptTemplate.from_template(general_instructions)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "77adf8bb-fb72-4afd-b768-3853c8fdba44",
   "metadata": {},
   "outputs": [],
   "source": [
    "branch = RunnableBranch(\n",
    "    (lambda x: \"support\" in x[\"topic\"].lower(), support_chain),\n",
    "    (lambda x: \"sales\" in x[\"topic\"].lower(), sales_chain),\n",
    "    (lambda x: \"joke\" in x[\"topic\"].lower(), joke_chain),\n",
    "    general_chain,\n",
    ") | RunnableLambda(lambda x: x.text)\n",
    "\n",
    "branch_chain = {\"topic\": classification_chain, \"human_input\": lambda x: x[\"human_input\"]} | branch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "a70cac41-3450-4241-9f8d-4f374ce40b9a",
   "metadata": {},
   "outputs": [],
   "source": [
    "response = branch_chain.invoke({\"human_input\": \"Can I track my order? I'm eager to know its status\"})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "15350c30-a240-41aa-94c4-6679606abd22",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\"You are a customer support agent. It seems that the user may have some issues. Answer to their query politely and sincerely. Be kind, understanding and say you're sorry for the inconvenience or the situation whenever necessary. Be brief and to the point.\""
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "response"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c80f2c19-22bb-43cb-b9d1-bfd6edf3541a",
   "metadata": {},
   "source": [
    "### Step 4: add full chat template and chain\n",
    "\n",
    "We finally define the final template with the different parts and bits we've defined above. It consists of three main sections:\n",
    "\n",
    "+ The starting line with some broad instructions on how to behave when responding in the chat.\n",
    "+ A `Instructions:` section where we'll place specific instructions for our three categories (`support`, `sales` and `joke`).\n",
    "+ The chat section, where we place our `chat_history` and the new `human_input`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "35e74216-ab95-4069-b994-c8f5064fc60e",
   "metadata": {},
   "outputs": [],
   "source": [
    "template = \"\"\"\\\n",
    "You are a chatbot having a conversation with a human. Follow the given instructions to reply to the Human message below.\n",
    "\n",
    "Instructions:{instructions}\n",
    "\n",
    "{chat_history}\n",
    "Human: {human_input}\n",
    "Chatbot:\"\"\"\n",
    "\n",
    "prompt = PromptTemplate(\n",
    "    input_variables=[\"instructions\", \"chat_history\", \"human_input\"], template=template\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "19b9d20f-8f09-409b-a8d0-6f72c4d49259",
   "metadata": {},
   "source": [
    "Out `chat_template` below just gives us the `prompt`. In other sections of the project, we've used chains that ended with a call to the language model (`chain = ... | hf`). Here we won't do that, though. The reason is explained on the next section."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "17150d76-bc1f-4c0f-b064-66afe42f6a90",
   "metadata": {},
   "outputs": [],
   "source": [
    "chat_chain = (\n",
    "    {\n",
    "        \"human_input\": lambda x: x[\"human_input\"], \n",
    "        \"instructions\": lambda x: branch_chain,\n",
    "        \"chat_history\": lambda x: x[\"chat_history\"],\n",
    "    } | prompt\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "589347c2-fedc-4910-82bc-8460e200b4b0",
   "metadata": {},
   "source": [
    "## Add user interface\n",
    "\n",
    "We use `gradio`'s `ChatInterface` to create a quick UI to use and test the chatbot. \n",
    "\n",
    "> You can set `DEBUG = True` below to see the prompts sent to the LLM,\n",
    "\n",
    "#### Explanation\n",
    "\n",
    "The `predict` function from `ChatInterface` gives us both the new user `message` and the `history` of messages. That's the reason why we don't need any kind of `langchain`'s memory here.\n",
    "\n",
    "In addition, we'll use a chain (`chat_chain`) that ends in a prompt, without making the call to the Huggingface model. The reason for that is that we were unable to find a way of defining a full langchain chain and making use of the `stream`ing feature so that the text would appear word by word. We researched how to do this, but in the end we couldn't make it work. [In some places](https://github.com/langchain-ai/langchain/issues/2918#issuecomment-1516441771) they recommend to use `HuggingFacePipeline`, but [it's not clear streaming is even supported by it](https://python.langchain.com/docs/integrations/llms/#features-natively-supported). Thus, we break the chain and divide it into two parts:\n",
    "\n",
    "+ part one, create the right prompt for the `human_input`\n",
    "+ part two, call the model using a _streamer_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "829a254b-db17-482f-907f-6a4eaced0eac",
   "metadata": {},
   "outputs": [],
   "source": [
    "import gradio as gr\n",
    "import torch\n",
    "from load_phi_model import StopOnTokens, StopOnNames\n",
    "from transformers import StoppingCriteriaList, TextIteratorStreamer, pipeline\n",
    "from threading import Thread\n",
    "import re"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "id": "1ca50386-e5c6-4dfa-831a-3d07e0e94dd6",
   "metadata": {},
   "outputs": [],
   "source": [
    "HUMAN_NAME = \"Human\"\n",
    "BOT_NAME = \"Chatbot\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "id": "9e125c85-5086-4577-b129-1702898ed661",
   "metadata": {},
   "outputs": [],
   "source": [
    "DEBUG = False # set to True to see the prompt sent to the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "id": "3ab4684d-0ff0-4833-9f4c-5987606ebcda",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Running on local URL:  http://127.0.0.1:7880\n",
      "\n",
      "To create a public link, set `share=True` in `launch()`.\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div><iframe src=\"http://127.0.0.1:7880/\" width=\"100%\" height=\"500\" allow=\"autoplay; camera; microphone; clipboard-read; clipboard-write;\" frameborder=\"0\" allowfullscreen></iframe></div>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": []
     },
     "execution_count": 50,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "chat_name_pattern_end = r'\\n.+:$' # matches substrings like `\\nUser:` at the end\n",
    "\n",
    "def predict(message, history):\n",
    "    stop_on_tokens = StopOnTokens()\n",
    "    stop_on_names = StopOnNames(\n",
    "        [tokenizer.encode(HUMAN_NAME), tokenizer.encode(BOT_NAME)])\n",
    "\n",
    "    messages = \"\".join([\"\".join(\n",
    "        [f\"\\n{HUMAN_NAME}: \"+item[0], f\"\\n{BOT_NAME}:\"+item[1]]\n",
    "    ) for item in history]).strip()\n",
    "\n",
    "    input_dict = {\n",
    "        \"human_input\": message,\n",
    "        \"chat_history\": messages,\n",
    "    }\n",
    "\n",
    "    input_prompt = chat_chain.invoke(input_dict).text\n",
    "    if DEBUG: print(input_prompt)\n",
    "\n",
    "    model_inputs = tokenizer([input_prompt], return_tensors=\"pt\").to(device)\n",
    "    streamer = TextIteratorStreamer(tokenizer, timeout=10., \n",
    "                                    skip_prompt=True, skip_special_tokens=True)\n",
    "    generate_kwargs = dict(\n",
    "        model_inputs,\n",
    "        streamer=streamer,\n",
    "        max_new_tokens=256,\n",
    "        do_sample=True,\n",
    "        top_p=0.95,\n",
    "        top_k=1000,\n",
    "        temperature=1.0,\n",
    "        num_beams=1,\n",
    "        stopping_criteria=StoppingCriteriaList([stop_on_tokens, stop_on_names])\n",
    "        )\n",
    "    t = Thread(target=model.generate, kwargs=generate_kwargs)\n",
    "    t.start()\n",
    "\n",
    "    partial_message = \"\"\n",
    "    for new_token in streamer:\n",
    "        partial_message += new_token\n",
    "        match = re.search(chat_name_pattern_end, partial_message)\n",
    "        if match:\n",
    "            partial_message = partial_message[:-len(match.group())]\n",
    "        yield partial_message\n",
    "        \n",
    "\n",
    "gr.ChatInterface(predict).queue().launch()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python3 Ipglobal 2",
   "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"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}