experiments_langchain.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a5e2844c-096a-45d8-8dab-e1b6ea5bf5b8",
   "metadata": {},
   "source": [
    "# Experiments with LangChain\n",
    "\n",
    "This notebook contains several code snippets and experiments done with [LangChain](https://www.langchain.com/)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "74ad14a7-80ba-48ce-adc9-7b14bd0b89e1",
   "metadata": {},
   "source": [
    "## Prompt for text classification\n",
    "\n",
    "The first thing we're going to do is to implement the same prompt we did in `experiments` for text classification, but using LangChain."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "edce4b90-edeb-461e-a753-a65ed61c834f",
   "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"
     ]
    }
   ],
   "source": [
    "from langchain.prompts import ChatPromptTemplate\n",
    "import torch\n",
    "from transformers import (\n",
    "    AutoModelForCausalLM, \n",
    "    AutoTokenizer, \n",
    "    StoppingCriteria, \n",
    "    StoppingCriteriaList, \n",
    "    TextIteratorStreamer\n",
    "    )\n",
    "from samples import samples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "734021fb-1cae-474a-a9cc-261e9b4370ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "template_string = \"\"\"\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: {text}\n",
    "Output:\n",
    "\"\"\".strip()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "2685fe86-a39a-4e6c-9fc5-e5b6ea2b0188",
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt_template = ChatPromptTemplate.from_template(template_string)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "d4bc1d61-0907-4703-bac3-e2c8de9f692f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Input variables are ['text']\n"
     ]
    }
   ],
   "source": [
    "print(\"Input variables are\", prompt_template.messages[0].input_variables)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "17d89414-e088-49f9-8970-132583a4c09b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "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: My package was lost\n",
      "Output:\n"
     ]
    }
   ],
   "source": [
    "print(prompt_template.format_messages(text=\"My package was lost\")[0].content)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "28807fb9-84a7-4a7a-9313-b6cb7dd5b2d4",
   "metadata": {},
   "source": [
    "## Connect `langchain` with `phi-2` and `huggingface`\n",
    "\n",
    "Ok, so now that we have a way to create prompts for our text classification using LangChain, we need to connect it with the AI model.\n",
    "\n",
    "Let's start with some already know code to load and use the `phi-2` model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "1a0e4327-8097-4c6f-ab59-004621e27e67",
   "metadata": {},
   "outputs": [
    {
     "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.43it/s]\n",
      "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n"
     ]
    }
   ],
   "source": [
    "class StopOnTokens(StoppingCriteria):\n",
    "    \"\"\"Stops the model if it produces an 'end of text' token\"\"\"\n",
    "    def __call__(self, input_ids: torch.LongTensor, \n",
    "                 scores: torch.FloatTensor, **kwargs) -> bool:\n",
    "        stop_ids = [50256, 198] # <|endoftext|> and EOL\n",
    "        for stop_id in stop_ids:\n",
    "            if input_ids[0][-1] == stop_id:\n",
    "                return True\n",
    "        return False\n",
    "\n",
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "print(\"Your device is\", device)\n",
    "\n",
    "model = AutoModelForCausalLM.from_pretrained(\n",
    "    \"microsoft/phi-2\", \n",
    "    device_map=\"auto\", \n",
    "    torch_dtype=\"auto\" if device == \"cuda\" else torch.float, \n",
    "    trust_remote_code=True\n",
    ")\n",
    "tokenizer = AutoTokenizer.from_pretrained(\"microsoft/phi-2\", trust_remote_code=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ee847aee-2461-440e-941c-009f1db3306b",
   "metadata": {},
   "source": [
    "To do that, we need to do it through `langchain`'s `HuggingFacePipeline`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "5e79cf61-7651-4bc6-a6b1-98c7701cd11b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.llms.huggingface_pipeline import HuggingFacePipeline\n",
    "from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline\n",
    "\n",
    "pipe = pipeline(\"text-generation\", model=model, tokenizer=tokenizer, max_new_tokens=200, stopping_criteria=[StopOnTokens()])\n",
    "hf = HuggingFacePipeline(pipeline=pipe)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cf0a76c5-9a2d-48e2-95ea-27b66482d4b9",
   "metadata": {},
   "source": [
    "Now define a new chain (with the `prompt_template` we defined in the previous section) and test it with some random sentence from our dataset. To do so, we use [LangChain Expression Language (LCEL)](https://python.langchain.com/docs/expression_language/), which is the new way of composing chains together."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "143e6d02-d8b3-49a3-a9fb-a1bcce36c86d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.output_parsers import StrOutputParser\n",
    "\n",
    "chain = prompt_template | hf | StrOutputParser()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "6e99bf06-2123-4a5d-ba5f-5f6d7564b048",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "PREDICTED: joke\n",
      "TRUE CATEGORY joke\n",
      "TEXT: How does a chatbot make a decision? It uses algorithm-ination!\n"
     ]
    }
   ],
   "source": [
    "from random import choice\n",
    "\n",
    "category = choice(list(samples.keys()))\n",
    "text = choice(samples[category])\n",
    "\n",
    "response = chain.invoke({\"text\": text})\n",
    "\n",
    "print(\"PREDICTED:\", response.strip())\n",
    "print(\"TRUE CATEGORY\", category)\n",
    "print(\"TEXT:\", text)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ec6cc807-3bb1-44f9-b61a-7ed5ada25c38",
   "metadata": {},
   "source": [
    "As we can see, our first `chain` is able to classify the input `text` in the following categories: support, sales and joke."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7b5c50e0-131f-4c9a-8ae7-05b74f2ae516",
   "metadata": {},
   "source": [
    "## Routing\n",
    "\n",
    "The next thing we have to do is to implement some routing to branch depending on the classification result. We want the chatbot to accomodate to the tone and content of the conversation. Thus, we'll set different instructions to follow depending on whether we're talking about _customer support_, _sales_, or simply _joking_.\n",
    "\n",
    "We're going to use `RunnableBranch` to run a different chain depending on the classification output. In our case we're going to implement a different prompt each time. There is also a generic prompt in the case the classifier is unable to undertand the content of the text or the output of the model is not properly parsed.\n",
    "\n",
    "Below, we define four different prompts and chains."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "e3ea4e86-b27c-4f25-bd00-4a7e9c303a56",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.prompts import PromptTemplate\n",
    "\n",
    "support_template = \"\"\"\\\n",
    "Instruction: 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.\n",
    "Query: {text}\n",
    "Output:\\\n",
    "\"\"\"\n",
    "\n",
    "sales_template = \"\"\"\\\n",
    "Instruction: 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",
    "Query: {text}\n",
    "Output:\\\n",
    "\"\"\"\n",
    "\n",
    "joke_template = \"\"\"\\\n",
    "Instruction: You are a comedian. The user want's to have some fun. Reply to their query in a funny way.\n",
    "Query: {text}\n",
    "Output:\\\n",
    "\"\"\"\n",
    "\n",
    "general_template = \"\"\"\\\n",
    "Instruction: Respond to the following query.\n",
    "Query: {text}\n",
    "Output:\\\n",
    "\"\"\"\n",
    "\n",
    "support_chain = PromptTemplate.from_template(support_template) | hf\n",
    "sales_chain = PromptTemplate.from_template(sales_template) | hf\n",
    "joke_chain = PromptTemplate.from_template(joke_template) | hf\n",
    "general_chain = PromptTemplate.from_template(general_template) | hf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "c1e23325-1562-42ac-b3b3-47ea3b1743fb",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.runnables import RunnableBranch\n",
    "\n",
    "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",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "7585bcdf-f5bb-4902-9bd0-c90834c7568e",
   "metadata": {},
   "outputs": [],
   "source": [
    "chain_2 = {\"topic\": chain, \"text\": lambda x: x[\"text\"]} | branch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "c5f772e2-5dc5-4aa2-a564-50c8f5afea18",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "QUERY How does a chatbot make a decision? It uses algorithm-ination!\n",
      "RESPONSE  Ha ha, that's hilarious! I love it! You're so clever and witty! You should be a comedian yourself!\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(\"QUERY\", text)\n",
    "print(\"RESPONSE\", chain_2.invoke({\"text\": text}))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "92416983-dc15-4416-8bdf-0eaa056c506f",
   "metadata": {},
   "source": [
    "## Add memory to the system\n",
    "\n",
    "Finally, we need to add memory to the system so that the model can remember past turns of the conversation. To do so we use `ConversationBufferMemory`.\n",
    "\n",
    "**Note**: We won't need this `memory` in the final code because our UI provides the list of saved messages for us."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "64da1b7b-3816-4502-96a2-9fd4ec95c730",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.chains import LLMChain\n",
    "from langchain.memory import ConversationBufferMemory\n",
    "from langchain.prompts import PromptTemplate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "13860f23-f623-4dd1-acb4-8faa855f3ec6",
   "metadata": {},
   "outputs": [],
   "source": [
    "template = \"\"\"You are a chatbot having a conversation with a human.\n",
    "\n",
    "{chat_history}\n",
    "Human: {human_input}\n",
    "Chatbot:\"\"\"\n",
    "\n",
    "prompt = PromptTemplate(\n",
    "    input_variables=[\"chat_history\", \"human_input\"], template=template\n",
    ")\n",
    "memory = ConversationBufferMemory(memory_key=\"chat_history\", input_key=\"human_input\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "ed4fc557-ab19-4eaf-9aaf-b96037e4719f",
   "metadata": {},
   "outputs": [],
   "source": [
    "llm_chain = LLMChain(\n",
    "    llm=hf,\n",
    "    prompt=prompt,\n",
    "    verbose=True,\n",
    "    memory=memory,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "b09564ac-f16c-48f8-ad6c-3c7bcc9679a6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new LLMChain chain...\u001b[0m\n",
      "Prompt after formatting:\n",
      "\u001b[32;1m\u001b[1;3mYou are a chatbot having a conversation with a human.\n",
      "\n",
      "\n",
      "Human: Hi there my friend\n",
      "Chatbot:\u001b[0m\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/grego/code/jobsearch/ipglobal-2/env/lib/python3.10/site-packages/transformers/pipelines/base.py:1101: UserWarning: You seem to be using the pipelines sequentially on GPU. In order to maximize efficiency please use a dataset\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "' Hello! How can I help you today?\\n'"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "llm_chain.predict(human_input=\"Hi there my friend\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "5b5e5aba-486c-496f-aa7d-4b8e0fc86621",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new LLMChain chain...\u001b[0m\n",
      "Prompt after formatting:\n",
      "\u001b[32;1m\u001b[1;3mYou are a chatbot having a conversation with a human.\n",
      "\n",
      "Human: Hi there my friend\n",
      "AI:  Hello! How can I help you today?\n",
      "\n",
      "Human: Can you give me some tips for studying?\n",
      "Chatbot:\u001b[0m\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/grego/code/jobsearch/ipglobal-2/env/lib/python3.10/site-packages/transformers/pipelines/base.py:1101: UserWarning: You seem to be using the pipelines sequentially on GPU. In order to maximize efficiency please use a dataset\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "\" Sure! I'd be happy to help. What subject are you studying?\\n\""
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "llm_chain.predict(human_input=\"Can you give me some tips for studying?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "84fa8fdd-a16f-43db-9c80-ffd4ec6b8a00",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new LLMChain chain...\u001b[0m\n",
      "Prompt after formatting:\n",
      "\u001b[32;1m\u001b[1;3mYou are a chatbot having a conversation with a human.\n",
      "\n",
      "Human: Hi there my friend\n",
      "AI:  Hello! How can I help you today?\n",
      "\n",
      "Human: Can you give me some tips for studying?\n",
      "AI:  Sure! I'd be happy to help. What subject are you studying?\n",
      "\n",
      "Human: Quantum Physiscs\n",
      "Chatbot:\u001b[0m\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/grego/code/jobsearch/ipglobal-2/env/lib/python3.10/site-packages/transformers/pipelines/base.py:1101: UserWarning: You seem to be using the pipelines sequentially on GPU. In order to maximize efficiency please use a dataset\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "\"  That's a great choice! Have you tried breaking down the material into smaller chunks and focusing on one topic at a time? It can help make the material more manageable. Also, don't forget to take breaks and give your brain a rest. Would you like me to help you create a study plan?\\n\""
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "llm_chain.predict(human_input=\"Quantum Physiscs\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9fd780f2-1e28-4816-9c7e-1cbfbc993281",
   "metadata": {},
   "source": [
    "## Conclusion\n",
    "\n",
    "In this notebook, we explore the mechanism that we need to implement the chatbot. Now, the only theing that's missing is to put it all together, but that's something we'll do in another notebook."
   ]
  }
 ],
 "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
}