00-mixtral-8x7b-agent.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "b0807f0c",
   "metadata": {},
   "source": [
    "# Mixtral 8x7B"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eb6ca53c",
   "metadata": {},
   "source": [
    "In this notebook we'll get started with Mistral AI's new model — **Mixtral 8x7B**. We'll see how to get setup with Mixtral, the prompt format it requires, and how it performs when being used as an Agent.\n",
    "\n",
    "As a bit of a spoiler, Mixtral is probably the first open source LLM that is truly _very very_ good — we say this considering the following key points:\n",
    "\n",
    "* Benchmarks show it to perform better than GPT-3.5.\n",
    "* Our own testing show Mixtral to be the first open source model we can reliably use as an agent.\n",
    "* Due to MoE architecture it is _very_ fast given it's size. If you can afford to run on 2x A100s latency is good enough to be used in chatbot use-cases.\n",
    "\n",
    "With that in mind, Mixtral is still 8x models — the total number of parameters is ~56B, so we do still need plenty of space to store the model. It's likely the amount of space required will decrease soon as quantized versions of the model are released (update, thanks [TheBloke](https://huggingface.co/TheBloke/Mixtral-8x7B-v0.1-GGUF))."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "46f331e9",
   "metadata": {},
   "source": [
    "## Finding Somewhere to Run Mixtral"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a7444bd6",
   "metadata": {},
   "source": [
    "Unless you have two A100s or H100s lying around you'll need to find a service to run Mixtral. We'll demonstrate how using [RunPod] TK here — we found this to be one of the easiest and cheapers compute providers to setup with Mixtral.\n",
    "\n",
    "* First, you'll need to [sign up for an account on RunPod] TK link.\n",
    "* Navigate to **Home** > click **Start Building**.\n",
    "* Setup a GPU instance, you can use 2xA100 or 2xH100.\n",
    "* Customize deployment to use *Container Size: 120GB* and *Disk Volume: 600GB*.\n",
    "* Make sure *Jupyter Notebook* is checked and deploy!\n",
    "\n",
    "Once deployed, click on the instance and click *Open Jupyter Server* — this will take you to a Jupyter Labs instance running on the container. From there you can upload _this notebook_ and follow along from there."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "64de2491",
   "metadata": {},
   "source": [
    "## Installing Prerequisites"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b24bad34",
   "metadata": {},
   "source": [
    "There are a few prerequisites required, to run Mixtral 8x7B we need `transformers` and `accelerate`. We also install `duckduckgo_search` to use in our agent testing later."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "f6b47f6a-b729-4992-9344-00beae46057b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\u001b[33m\n",
      "\u001b[0m"
     ]
    }
   ],
   "source": [
    "%pip install -qU \\\n",
    "    transformers==4.36.1 \\\n",
    "    accelerate==0.25.0 \\\n",
    "    duckduckgo_search==4.1.0"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a3a7f024",
   "metadata": {},
   "source": [
    "_**IMPORTANT: You may need to restart the kernel before continuing for the above library installs to be recognized by the remainder of the code!**_"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "06099840",
   "metadata": {},
   "source": [
    "## Download and Initialize Mixtral"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "be819054-3c82-4f9f-a2b5-fae0f76f9033",
   "metadata": {},
   "source": [
    "After installing above, refresh the kernel before continuing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "fde21500-7439-4fd4-beef-0773b2d5b902",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "52213ff8bbd24bbabd6f0e2428ce42b9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/19 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "MixtralForCausalLM(\n",
       "  (model): MixtralModel(\n",
       "    (embed_tokens): Embedding(32000, 4096)\n",
       "    (layers): ModuleList(\n",
       "      (0-31): 32 x MixtralDecoderLayer(\n",
       "        (self_attn): MixtralAttention(\n",
       "          (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (k_proj): Linear(in_features=4096, out_features=1024, bias=False)\n",
       "          (v_proj): Linear(in_features=4096, out_features=1024, bias=False)\n",
       "          (o_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (rotary_emb): MixtralRotaryEmbedding()\n",
       "        )\n",
       "        (block_sparse_moe): MixtralSparseMoeBlock(\n",
       "          (gate): Linear(in_features=4096, out_features=8, bias=False)\n",
       "          (experts): ModuleList(\n",
       "            (0-7): 8 x MixtralBLockSparseTop2MLP(\n",
       "              (w1): Linear(in_features=4096, out_features=14336, bias=False)\n",
       "              (w2): Linear(in_features=14336, out_features=4096, bias=False)\n",
       "              (w3): Linear(in_features=4096, out_features=14336, bias=False)\n",
       "              (act_fn): SiLU()\n",
       "            )\n",
       "          )\n",
       "        )\n",
       "        (input_layernorm): MixtralRMSNorm()\n",
       "        (post_attention_layernorm): MixtralRMSNorm()\n",
       "      )\n",
       "    )\n",
       "    (norm): MixtralRMSNorm()\n",
       "  )\n",
       "  (lm_head): Linear(in_features=4096, out_features=32000, bias=False)\n",
       ")"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from torch import bfloat16\n",
    "import transformers\n",
    "\n",
    "model_id = \"mistralai/Mixtral-8x7B-Instruct-v0.1\"\n",
    "\n",
    "model = transformers.AutoModelForCausalLM.from_pretrained(\n",
    "    model_id,\n",
    "    trust_remote_code=True,\n",
    "    torch_dtype=bfloat16,\n",
    "    device_map='auto'\n",
    ")\n",
    "model.eval()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "38bffa2d",
   "metadata": {},
   "source": [
    "As with all LLMs/transformer models we need to initialize a `tokenizer` that will take our plaintext inputs and transform them into lists of tokens that are consumed by the first layer of the LLM/transformer. We initialize the Mixtral tokenizer like so:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "2ba29e2c-a666-4767-971d-ad2ddbb2dc86",
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer = transformers.AutoTokenizer.from_pretrained(model_id)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "36a3babd",
   "metadata": {},
   "source": [
    "Now we setup a `text-generation` pipeline using `transformers`. There are a lot of generation parameters we can adjust here, we'd recommend leaving them as is for now and returning to them if you feel like your generated outputs need improvement."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "7c62ed82-5de1-4827-a24e-05e55d8f6c57",
   "metadata": {},
   "outputs": [],
   "source": [
    "generate_text = transformers.pipeline(\n",
    "    model=model, tokenizer=tokenizer,\n",
    "    return_full_text=False,  # if using langchain set True\n",
    "    task=\"text-generation\",\n",
    "    # we pass model parameters here too\n",
    "    temperature=0.1,  # 'randomness' of outputs, 0.0 is the min and 1.0 the max\n",
    "    top_p=0.15,  # select from top tokens whose probability add up to 15%\n",
    "    top_k=0,  # select from top 0 tokens (because zero, relies on top_p)\n",
    "    max_new_tokens=512,  # max number of tokens to generate in the output\n",
    "    repetition_penalty=1.1  # if output begins repeating increase\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "000877a4",
   "metadata": {},
   "source": [
    "Now we can generate text by calling `generate_text`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "e0069d13-8951-4970-8387-997a88cdbf4b",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.10/dist-packages/transformers/generation/configuration_utils.py:389: UserWarning: `do_sample` is set to `False`. However, `temperature` is set to `0.1` -- this flag is only used in sample-based generation modes. You should set `do_sample=True` or unset `temperature`.\n",
      "  warnings.warn(\n",
      "/usr/local/lib/python3.10/dist-packages/transformers/generation/configuration_utils.py:394: UserWarning: `do_sample` is set to `False`. However, `top_p` is set to `0.15` -- this flag is only used in sample-based generation modes. You should set `do_sample=True` or unset `top_p`.\n",
      "  warnings.warn(\n",
      "/usr/local/lib/python3.10/dist-packages/transformers/generation/configuration_utils.py:404: UserWarning: `do_sample` is set to `False`. However, `top_k` is set to `0` -- this flag is only used in sample-based generation modes. You should set `do_sample=True` or unset `top_k`.\n",
      "  warnings.warn(\n",
      "Setting `pad_token_id` to `eos_token_id`:2 for open-end generation.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      ", #TeamInternet!\n",
      "\n",
      "I’m so excited to be here today as a guest blogger for the lovely ladies of The Blogcademy. I’ve been following their adventures since they first started teaching classes in London and New York City last year, and it was such an honor when Kat asked me if I would like to write a post about my experience with social media.\n",
      "\n",
      "As you may know, I am the founder of TweetHearts, a social media management company that specializes in Twitter marketing. We help businesses grow their online presence by creating customized strategies that are tailored to meet their specific needs. Our clients range from small start-ups to large corporations, but no matter what size or industry, we always focus on building relationships through authentic engagement.\n",
      "\n",
      "When it comes to social media, many people think that all you have to do is set up an account and then just start posting updates whenever you feel like it. However, this couldn’t be further from the truth! In order to truly succeed in the world of social media, you need to have a solid plan in place – one that includes clear goals, targeted messaging, consistent branding, and most importantly, a commitment to engaging with your audience.\n",
      "\n",
      "Here are some tips to help get you started:\n",
      "\n",
      "1. Define Your Goals\n",
      "Before you even create your social media profiles, take some time to think about what you hope to achieve through these platforms. Are you looking to increase brand awareness? Drive traffic to your website? Build relationships with potential customers? Whatever your goals may be, make sure they are specific, measurable, attainable, relevant, and timely (SMART). This will not only give you a clear direction for your social media efforts, but it will also allow you to track your progress and adjust your strategy accordingly.\n",
      "\n",
      "2. Know Your Audience\n",
      "In order to effectively engage with your audience, you need to understand who they are, what they care about, and how they use social media. Take a look at your current customer base and identify any common demographics, interests, or behaviors. Then, research which social media platforms they are most active on and what types of content resonate with them. By doing this, you can create targeted messages that speak directly to your audience and encourage them to interact with your brand.\n",
      "\n",
      "3. Create a Content Strategy\n",
      "Once you know who your audience is and what they want to see, it’s\n"
     ]
    }
   ],
   "source": [
    "res = generate_text(\"hello there\")\n",
    "print(res[0][\"generated_text\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b0aad106",
   "metadata": {},
   "source": [
    "### Instruction Format"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "626b9fe8",
   "metadata": {},
   "source": [
    "We can see a very generic generated output here. There are two primary reasons for that:\n",
    "\n",
    "1. We haven't provided any instructions to the model.\n",
    "2. We have not used the recommended instruction format.\n",
    "\n",
    "The instruction format for Mixtral 8x7B looks like this:\n",
    "\n",
    "```\n",
    "<s> [INST] Some instructions [/INST] Primer text [generated output] </s>\n",
    "```\n",
    "\n",
    "We would put our instructions to the model in place of `\"Some instructions\"` and place a primer like `\"Assistant: \"` in place of `\"Primer text\"`. The `<s>` and `</s>` are special tokens used by Mixtral to signify the **B**eginning **O**f **S**tring (BOS) and **E**nd **O**f **S**tring (EOS), ie beginning and end of our text. The `[INST]` and `[/INST]` strings tell the model that anything between those two strings are _instructions_ that the model should follow.\n",
    "\n",
    "We can add some follow up instructions like so:\n",
    "\n",
    "```\n",
    "<s> [INST] Some instructions [/INST] Primer text [generated output] </s> [INST] Further instructions [/INST]\n",
    "```\n",
    "\n",
    "Let's begin by adding some _instructions_ first, we'll add instruction formatting later. In these instructions we want to setup the guidelines for an agent that can use two tools (calculator and search) and also return answers to the user. All three of these options will be used by the agent via a JSON output format containing `\"tool_name\"` that specifies which tool to be used (one of [`Calculator`, `Search`, `Final Answer`]) and `\"input\"` that specifies the input to the chosen tool."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "9c4f2154-39d9-412e-9a0c-b0e935671ab7",
   "metadata": {},
   "outputs": [],
   "source": [
    "agent_template = \"\"\"\n",
    "You are a helpful AI assistant, you are an agent capable of using a variety of tools to answer a question. Here are a few of the tools available to you:\n",
    "\n",
    "- Calculator: the calculator should be used whenever you need to perform a calculation, no matter how simple. It uses Python so make sure to write complete Python code required to perform the calculation required and make sure the Python returns your answer to the `output` variable.\n",
    "- Search: the search tool should be used whenever you need to find information. It can be used to find information about everything\n",
    "- Final Answer: the final answer tool must be used to respond to the user. You must use this when you have decided on an answer.\n",
    "\n",
    "To use these tools you must always respond in JSON format containing `\"tool_name\"` and `\"input\"` key-value pairs. For example, to answer the question, \"what is the square root of 51?\" you must use the calculator tool like so:\n",
    "\n",
    "```json\n",
    "{\n",
    "    \"tool_name\": \"Calculator\",\n",
    "    \"input\": \"from math import sqrt; output = sqrt(51)\"\n",
    "}\n",
    "```\n",
    "\n",
    "Or to answer the question \"who is the current president of the USA?\" you must respond:\n",
    "\n",
    "```json\n",
    "{\n",
    "    \"tool_name\": \"Search\",\n",
    "    \"input\": \"current president of USA\"\n",
    "}\n",
    "```\n",
    "\n",
    "Remember, even when answering to the user, you must still use this JSON format! If you'd like to ask how the user is doing you must write:\n",
    "\n",
    "```json\n",
    "{\n",
    "    \"tool_name\": \"Final Answer\",\n",
    "    \"input\": \"How are you today?\"\n",
    "}\n",
    "```\n",
    "\n",
    "Let's get started. The users query is as follows.\n",
    "\n",
    "User: Hi there, I'm stuck on a math problem, can you help? My question is what is the square root of 512 multiplied by 7?\n",
    "\n",
    "Assistant: ```json\n",
    "{\n",
    "    \"tool_name\": \"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "57f3ec8f",
   "metadata": {},
   "source": [
    "Using these instructions alone we actually get great performance:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "8c2cc444-fd93-43c0-b658-e5fc78e5e03a",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Setting `pad_token_id` to `eos_token_id`:2 for open-end generation.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " \"Calculator\",\n",
      "    \"input\":  \"from math import sqrt; output = sqrt(512) * 7\"\n",
      "}\n",
      "```\n"
     ]
    }
   ],
   "source": [
    "res = generate_text(agent_template)\n",
    "print(res[0][\"generated_text\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eca74dfe",
   "metadata": {},
   "source": [
    "But before continuing let's add the instruction formatting too, we'll do this via a function called `instruction_format` that will consume a `sys_message` (ie instructions) and a user's `query` and output the string with the required tokens."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "95e4da2b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def instruction_format(sys_message: str, query: str):\n",
    "    # note, don't \"</s>\" to the end\n",
    "    return f'<s> [INST] {sys_message} [/INST]\\nUser: {query}\\nAssistant: ```json\\n{{\\n\"tool_name\": '"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "02b589ec",
   "metadata": {},
   "source": [
    "Let's see what this looks like:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "ae2feb1b",
   "metadata": {},
   "outputs": [],
   "source": [
    "sys_msg = \"\"\"You are a helpful AI assistant, you are an agent capable of using a variety of tools to answer a question. Here are a few of the tools available to you:\n",
    "\n",
    "- Calculator: the calculator should be used whenever you need to perform a calculation, no matter how simple. It uses Python so make sure to write complete Python code required to perform the calculation required and make sure the Python returns your answer to the `output` variable.\n",
    "- Search: the search tool should be used whenever you need to find information. It can be used to find information about everything\n",
    "- Final Answer: the final answer tool must be used to respond to the user. You must use this when you have decided on an answer.\n",
    "\n",
    "To use these tools you must always respond in JSON format containing `\"tool_name\"` and `\"input\"` key-value pairs. For example, to answer the question, \"what is the square root of 51?\" you must use the calculator tool like so:\n",
    "\n",
    "```json\n",
    "{\n",
    "    \"tool_name\": \"Calculator\",\n",
    "    \"input\": \"from math import sqrt; output = sqrt(51)\"\n",
    "}\n",
    "```\n",
    "\n",
    "Or to answer the question \"who is the current president of the USA?\" you must respond:\n",
    "\n",
    "```json\n",
    "{\n",
    "    \"tool_name\": \"Search\",\n",
    "    \"input\": \"current president of USA\"\n",
    "}\n",
    "```\n",
    "\n",
    "Remember, even when answering to the user, you must still use this JSON format! If you'd like to ask how the user is doing you must write:\n",
    "\n",
    "```json\n",
    "{\n",
    "    \"tool_name\": \"Final Answer\",\n",
    "    \"input\": \"How are you today?\"\n",
    "}\n",
    "```\n",
    "\n",
    "Let's get started. The users query is as follows.\n",
    "\"\"\"\n",
    "query = \"Hi there, I'm stuck on a math problem, can you help? My question is what is the square root of 512 multiplied by 7?\"\n",
    "\n",
    "input_prompt = instruction_format(sys_msg, query)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "fa6db197-7bde-4c25-9406-c9b0e126c401",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<s> [INST] You are a helpful AI assistant, you are an agent capable of using a variety of tools to answer a question. Here are a few of the tools available to you:\n",
      "\n",
      "- Calculator: the calculator should be used whenever you need to perform a calculation, no matter how simple. It uses Python so make sure to write complete Python code required to perform the calculation required and make sure the Python returns your answer to the `output` variable.\n",
      "- Search: the search tool should be used whenever you need to find information. It can be used to find information about everything\n",
      "- Final Answer: the final answer tool must be used to respond to the user. You must use this when you have decided on an answer.\n",
      "\n",
      "To use these tools you must always respond in JSON format containing `\"tool_name\"` and `\"input\"` key-value pairs. For example, to answer the question, \"what is the square root of 51?\" you must use the calculator tool like so:\n",
      "\n",
      "```json\n",
      "{\n",
      "    \"tool_name\": \"Calculator\",\n",
      "    \"input\": \"from math import sqrt; output = sqrt(51)\"\n",
      "}\n",
      "```\n",
      "\n",
      "Or to answer the question \"who is the current president of the USA?\" you must respond:\n",
      "\n",
      "```json\n",
      "{\n",
      "    \"tool_name\": \"Search\",\n",
      "    \"input\": \"current president of USA\"\n",
      "}\n",
      "```\n",
      "\n",
      "Remember, even when answering to the user, you must still use this JSON format! If you'd like to ask how the user is doing you must write:\n",
      "\n",
      "```json\n",
      "{\n",
      "    \"tool_name\": \"Final Answer\",\n",
      "    \"input\": \"How are you today?\"\n",
      "}\n",
      "```\n",
      "\n",
      "Let's get started. The users query is as follows.\n",
      " [/INST]\n",
      "User: Hi there, I'm stuck on a math problem, can you help? My question is what is the square root of 512 multiplied by 7?\n",
      "Assistant: ```json\n",
      "{\n",
      "\"tool_name\": \n"
     ]
    }
   ],
   "source": [
    "print(input_prompt)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b6b49e0e",
   "metadata": {},
   "source": [
    "Now let's generate output from our LLM again:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "c8e4a97f",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Setting `pad_token_id` to `eos_token_id`:2 for open-end generation.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " \"Calculator\",\n",
      "\"input\":  \"from math import sqrt; output = sqrt(512) * 7\"\n",
      "}\n",
      "```\n"
     ]
    }
   ],
   "source": [
    "res = generate_text(input_prompt)\n",
    "print(res[0][\"generated_text\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bb338d39-9d1d-4f41-881b-05c50d0a46bf",
   "metadata": {},
   "source": [
    "Let's parse this into Python executable code..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "74458302-9900-40c1-a127-3098417685ee",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "def format_output(text: str):\n",
    "    full_json_str = '{\\n\"tool_name\": '+text\n",
    "    full_json_str = full_json_str.strip()\n",
    "    if full_json_str.endswith(\"```\"):\n",
    "        full_json_str = full_json_str[:-3]\n",
    "    return json.loads(full_json_str)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "aa73c53b-f5a6-43fb-a8a2-3611210560a3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'tool_name': 'Calculator',\n",
       " 'input': 'from math import sqrt; output = sqrt(512) * 7'}"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "action = format_output(res[0][\"generated_text\"])\n",
    "action"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "2bf74096-cd62-4291-96c5-f9647cd5b03a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "158.39191898578665"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "if action[\"tool_name\"] == \"Calculator\":\n",
    "    exec(action[\"input\"])\n",
    "output"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "608bd92b-f5c8-4cbd-b63f-1a24a4050f24",
   "metadata": {},
   "source": [
    "Now we add this info to our prompt:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "769b2bab-4da9-4ac0-a61b-064b2567568a",
   "metadata": {},
   "outputs": [],
   "source": [
    "new_prompt = f\"\"\"\n",
    "{agent_template}{res[0][\"generated_text\"]}\n",
    "Tool Output: {output}\n",
    "Assistant: ```json\n",
    "{{\n",
    "    \"tool_name\": \"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "1a3fd83d-8314-4672-a8df-b2a780d26b62",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "You are a helpful AI assistant, you are an agent capable of using a variety of tools to answer a question. Here are a few of the tools available to you:\n",
      "\n",
      "- Calculator: the calculator should be used whenever you need to perform a calculation, no matter how simple. It uses Python so make sure to write complete Python code required to perform the calculation required and make sure the Python returns your answer to the `output` variable.\n",
      "- Search: the search tool should be used whenever you need to find information. It can be used to find information about everything\n",
      "- Final Answer: the final answer tool must be used to respond to the user. You must use this when you have decided on an answer.\n",
      "\n",
      "To use these tools you must always respond in JSON format containing `\"tool_name\"` and `\"input\"` key-value pairs. For example, to answer the question, \"what is the square root of 51?\" you must use the calculator tool like so:\n",
      "\n",
      "```json\n",
      "{\n",
      "    \"tool_name\": \"Calculator\",\n",
      "    \"input\": \"from math import sqrt; output = sqrt(51)\"\n",
      "}\n",
      "```\n",
      "\n",
      "Or to answer the question \"who is the current president of the USA?\" you must respond:\n",
      "\n",
      "```json\n",
      "{\n",
      "    \"tool_name\": \"Search\",\n",
      "    \"input\": \"current president of USA\"\n",
      "}\n",
      "```\n",
      "\n",
      "Remember, even when answering to the user, you must still use this JSON format! If you'd like to ask how the user is doing you must write:\n",
      "\n",
      "```json\n",
      "{\n",
      "    \"tool_name\": \"Final Answer\",\n",
      "    \"input\": \"How are you today?\"\n",
      "}\n",
      "```\n",
      "\n",
      "Let's get started. The users query is as follows.\n",
      "\n",
      "User: Hi there, I'm stuck on a math problem, can you help? My question is what is the square root of 512 multiplied by 7?\n",
      "\n",
      "Assistant: ```json\n",
      "{\n",
      "    \"tool_name\":  \"Calculator\",\n",
      "\"input\":  \"from math import sqrt; output = sqrt(512) * 7\"\n",
      "}\n",
      "```\n",
      "Tool Output: 158.39191898578665\n",
      "Assistant: ```json\n",
      "{\n",
      "    \"tool_name\": \n"
     ]
    }
   ],
   "source": [
    "print(new_prompt)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "881bef62-0a93-4aff-88c7-1ee163baee65",
   "metadata": {},
   "source": [
    "And return this info back to the agent:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "c43b292a-11b7-4912-80b1-e92098dde0ed",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Setting `pad_token_id` to `eos_token_id`:2 for open-end generation.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " \"Final Answer\",\n",
      "\"input\":  \"The square root of 512 multiplied by 7 is approximately 158.4.\"\n",
      "}\n",
      "```\n"
     ]
    }
   ],
   "source": [
    "res2 = generate_text(new_prompt)\n",
    "print(res2[0][\"generated_text\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "69b12453-c048-4aca-b8ea-8f32ae7ad5c1",
   "metadata": {},
   "source": [
    "Convert to dictionary:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "dab28631-8553-4165-afe0-3681d6d11268",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'tool_name': 'Final Answer',\n",
       " 'input': 'The square root of 512 multiplied by 7 is approximately 158.4.'}"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "response = format_output(res2[0][\"generated_text\"])\n",
    "response"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8cd88c71-3fed-4975-9ad0-b401bbabeea8",
   "metadata": {},
   "source": [
    "We add some handling for `Final Answer` outputs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "4bef8dcd-8799-4d7d-9655-a5ec0d9c8fac",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Assistant: The square root of 512 multiplied by 7 is approximately 158.4.\n"
     ]
    }
   ],
   "source": [
    "if response[\"tool_name\"] == \"Final Answer\":\n",
    "    print(\"Assistant: \"+response[\"input\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b5566cf3-f401-4432-9750-961b563fd72b",
   "metadata": {},
   "source": [
    "Let's integrate this tool input parsing logging into a single function..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "fb20c9f0-2ff9-4da7-8da7-75904cf1ae0d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from duckduckgo_search import DDGS\n",
    "\n",
    "def use_tool(action: dict):\n",
    "    tool_name = action[\"tool_name\"]\n",
    "    if tool_name == \"Final Answer\":\n",
    "        return \"Assistant: \"+action[\"input\"]\n",
    "    elif tool_name == \"Calculator\":\n",
    "        exec(action[\"input\"])\n",
    "        return f\"Tool Output: {output}\"\n",
    "    elif tool_name == \"Search\":\n",
    "        contexts = []\n",
    "        with DDGS() as ddgs:\n",
    "            results = ddgs.text(\n",
    "                action[\"input\"],\n",
    "                region=\"wt-wt\", safesearch=\"on\",\n",
    "                max_results=3\n",
    "            )\n",
    "            for r in results:\n",
    "                contexts.append(r['body'])\n",
    "        info = \"\\n---\\n\".join(contexts)\n",
    "        return f\"Tool Output: {info}\"\n",
    "    else:\n",
    "        # otherwise just assume final answer\n",
    "        return \"Assistant: \"+action[\"input\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "5c41f2f1-a339-4709-9c15-c43561cc3e4e",
   "metadata": {},
   "outputs": [],
   "source": [
    "query = \"who is the current prime minister of the UK?\"\n",
    "\n",
    "input_prompt = instruction_format(sys_msg, query)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "f9eec9ed-5538-4a83-b363-4b145f550e29",
   "metadata": {},
   "outputs": [],
   "source": [
    "def run(query: str):\n",
    "    res = generate_text(query)\n",
    "    action_dict = format_output(res[0][\"generated_text\"])\n",
    "    response = use_tool(action_dict)\n",
    "    full_text = f\"{query}{res[0]['generated_text']}\\n{response}\"\n",
    "    return response, full_text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "eb91b472-e59d-41e3-8ecb-6c270788c75e",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Setting `pad_token_id` to `eos_token_id`:2 for open-end generation.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'Tool Output: Rishi Sunak became Prime Minister on 25 October 2022. He was previously appointed Chancellor of the Exchequer from 13 February 2020 to 5 July 2022.\\n---\\nRishi Sunak Political positions Electoral history MP for Richmond (Yorks) Prime Minister of the United Kingdom Premiership Minister for the Union Ministry 2023 February reshuffle 2023 November reshuffle Industrial disputes postal workers strikes railway strikes NHS strikes 2022 autumn statement Russian invasion of Ukraine economic impact\\n---\\nRishi Sunak has been the prime minister since 25 October 2022. [7] History Sir Robert Walpole is generally considered to have been the first person to hold the position of Prime Minister.'"
      ]
     },
     "execution_count": 39,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "out = run(input_prompt)\n",
    "out[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "e4cf726c-278c-4493-929f-5a5b101516c3",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Tool Output: Rishi Sunak became Prime Minister on 25 October 2022. He was previously appointed Chancellor of the Exchequer from 13 February 2020 to 5 July 2022.\n",
      "---\n",
      "Rishi Sunak Political positions Electoral history MP for Richmond (Yorks) Prime Minister of the United Kingdom Premiership Minister for the Union Ministry 2023 February reshuffle 2023 November reshuffle Industrial disputes postal workers strikes railway strikes NHS strikes 2022 autumn statement Russian invasion of Ukraine economic impact\n",
      "---\n",
      "Rishi Sunak has been the prime minister since 25 October 2022. [7] History Sir Robert Walpole is generally considered to have been the first person to hold the position of Prime Minister.\n"
     ]
    }
   ],
   "source": [
    "print(out[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "4f3e743a-8d68-488d-a90d-99ec56ae8102",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<s> [INST] You are a helpful AI assistant, you are an agent capable of using a variety of tools to answer a question. Here are a few of the tools available to you:\n",
      "\n",
      "- Calculator: the calculator should be used whenever you need to perform a calculation, no matter how simple. It uses Python so make sure to write complete Python code required to perform the calculation required and make sure the Python returns your answer to the `output` variable.\n",
      "- Search: the search tool should be used whenever you need to find information. It can be used to find information about everything\n",
      "- Final Answer: the final answer tool must be used to respond to the user. You must use this when you have decided on an answer.\n",
      "\n",
      "To use these tools you must always respond in JSON format containing `\"tool_name\"` and `\"input\"` key-value pairs. For example, to answer the question, \"what is the square root of 51?\" you must use the calculator tool like so:\n",
      "\n",
      "```json\n",
      "{\n",
      "    \"tool_name\": \"Calculator\",\n",
      "    \"input\": \"from math import sqrt; output = sqrt(51)\"\n",
      "}\n",
      "```\n",
      "\n",
      "Or to answer the question \"who is the current president of the USA?\" you must respond:\n",
      "\n",
      "```json\n",
      "{\n",
      "    \"tool_name\": \"Search\",\n",
      "    \"input\": \"current president of USA\"\n",
      "}\n",
      "```\n",
      "\n",
      "Remember, even when answering to the user, you must still use this JSON format! If you'd like to ask how the user is doing you must write:\n",
      "\n",
      "```json\n",
      "{\n",
      "    \"tool_name\": \"Final Answer\",\n",
      "    \"input\": \"How are you today?\"\n",
      "}\n",
      "```\n",
      "\n",
      "Let's get started. The users query is as follows.\n",
      " [/INST]\n",
      "User: who is the current prime minister of the UK?\n",
      "Assistant: ```json\n",
      "{\n",
      "\"tool_name\":  \"Search\",\n",
      "\"input\":  \"current prime minister of the UK\"\n",
      "}\n",
      "```\n",
      "Tool Output: Rishi Sunak became Prime Minister on 25 October 2022. He was previously appointed Chancellor of the Exchequer from 13 February 2020 to 5 July 2022.\n",
      "---\n",
      "Rishi Sunak Political positions Electoral history MP for Richmond (Yorks) Prime Minister of the United Kingdom Premiership Minister for the Union Ministry 2023 February reshuffle 2023 November reshuffle Industrial disputes postal workers strikes railway strikes NHS strikes 2022 autumn statement Russian invasion of Ukraine economic impact\n",
      "---\n",
      "Rishi Sunak has been the prime minister since 25 October 2022. [7] History Sir Robert Walpole is generally considered to have been the first person to hold the position of Prime Minister.\n"
     ]
    }
   ],
   "source": [
    "print(out[1])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d1d066b6-1d88-4509-88b4-eed8af9f5295",
   "metadata": {},
   "source": [
    "We're not handling the logic of iterating through multiple agent steps yet — so we need to do that manually..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "9ad9b8fc-bd7e-4849-9287-4e9bbb22e4e4",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.10/dist-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",
      "Setting `pad_token_id` to `eos_token_id`:2 for open-end generation.\n"
     ]
    }
   ],
   "source": [
    "second_step = out[1]+\"\"\"\n",
    "Assistant: ```json\n",
    "{\n",
    "    \"tool_name\": \"\"\"\n",
    "\n",
    "out = run(second_step)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "3d181518-2382-49ac-abb8-c0846dba5536",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'Assistant: Rishi Sunak has been the Prime Minister of the United Kingdom since 25 October 2022.'"
      ]
     },
     "execution_count": 43,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "out[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cfe7f34c-69ef-40f9-9c02-1d23d15fe0f1",
   "metadata": {},
   "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.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}