{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Prompt Engineering with Mixtral 8x7B\n",
    "\n",
    "This guide provides some prompt examples demonstrating how to use Mixtral 8x7B and it's wide range of capabilities. \n",
    "\n",
    "We will be using the official Python client from here: https://github.com/mistralai/client-python\n",
    "\n",
    "Make sure to setup a `MISTRAL_API_KEY` before getting started with the guide. You can it here: https://console.mistral.ai/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture\n",
    "!pip install mistralai"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Basic Usage"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "from mistralai.client import MistralClient\n",
    "from mistralai.models.chat_completion import ChatMessage\n",
    "from dotenv import load_dotenv\n",
    "\n",
    "load_dotenv()\n",
    "import os\n",
    "\n",
    "api_key = os.environ[\"MISTRAL_API_KEY\"]\n",
    "client = MistralClient(api_key=api_key)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "metadata": {},
   "outputs": [],
   "source": [
    "# helpful completion function\n",
    "\n",
    "def get_completion(messages, model=\"mistral-small\"):\n",
    "    # No streaming\n",
    "    chat_response = client.chat(\n",
    "        model=model,\n",
    "        messages=messages,\n",
    "    )\n",
    "\n",
    "    return chat_response\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "id='2c7bed47132b45ce8a76ec8e2c5df25d' object='chat.completion' created=1706504806 model='mistral-small' choices=[ChatCompletionResponseChoice(index=0, message=ChatMessage(role='assistant', content='Why do sharks swim in salt water? Because pepper would make them sneeze!'), finish_reason=<FinishReason.stop: 'stop'>)] usage=UsageInfo(prompt_tokens=15, total_tokens=32, completion_tokens=17)\n"
     ]
    }
   ],
   "source": [
    "messages = [\n",
    "    ChatMessage(role=\"user\", content=\"Tell me a joke about sharks\")\n",
    "]\n",
    "\n",
    "chat_response = get_completion(messages)\n",
    "print(chat_response)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'Why do sharks swim in salt water? Because pepper would make them sneeze!'"
      ]
     },
     "execution_count": 63,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# print only the content\n",
    "chat_response.choices[0].message.content"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Using the Chat Template\n",
    "\n",
    "To effectively prompt the Mistral 8x7B Instruct and get optimal outputs, it's recommended to use the following chat template:\n",
    "\n",
    "```\n",
    "<s>[INST] Instruction [/INST] Model answer</s>[INST] Follow-up instruction [/INST]\n",
    "```\n",
    "\n",
    "*Note that `<s>` and `</s>` are special tokens for beginning of string (BOS) and end of string (EOS) while [INST] and [/INST] are regular strings.*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 87,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{\n",
      "\"name\": \"John\",\n",
      "\"lastname\": \"Smith\",\n",
      "\"address\": \"#1 Samuel St.\"\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "prompt = \"\"\"[INST] You are a helpful code assistant. Your task is to generate a valid JSON object based on the given information:\n",
    "\n",
    "name: John\n",
    "lastname: Smith\n",
    "address: #1 Samuel St.\n",
    "\n",
    "Just generate the JSON object without explanations:\n",
    "[/INST]\"\"\"\n",
    "\n",
    "messages = [\n",
    "    ChatMessage(role=\"user\", content=prompt)\n",
    "]\n",
    "\n",
    "chat_response = get_completion(messages)\n",
    "print(chat_response.choices[0].message.content)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note the importance of the template that was used above. If we don't use the template, we get very different results. If we want to leverage the model capabilities in the proper way, we need to follow the format."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here is another example that uses a conversation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 83,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "My apologies, I said \"the right amount of zesty flavor.\" Fresh lemon juice can add a bright and tangy taste to various dishes, elevating their overall flavor profile.\n"
     ]
    }
   ],
   "source": [
    "prompt = \"\"\"<s>[INST] What is your favorite condiment? [/INST]\n",
    "\"Well, I'm quite partial to a good squeeze of fresh lemon juice. It adds just the right amount of zesty flavour to whatever I'm cooking up in the kitchen!\"</s> [INST] The right amount of what? [/INST]\"\"\"\n",
    "\n",
    "messages = [\n",
    "    ChatMessage(role=\"user\", content=prompt)\n",
    "]\n",
    "\n",
    "chat_response = get_completion(messages)\n",
    "print(chat_response.choices[0].message.content)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We could also use the `ChatMessage` to define the different roles and content."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The example below shows a similar task in a multi-turn conversation:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 85,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "My apologies for any confusion. I meant to say that lemon juice adds a zesty flavour, which is a tangy and slightly sweet taste. It's a delightful addition to many dishes, in my humble opinion.\n"
     ]
    }
   ],
   "source": [
    "messages = [\n",
    "    ChatMessage(role=\"user\", content=\"What is your favorite condiment?\"),\n",
    "    ChatMessage(role=\"assistant\", content=\"Well, I'm quite partial to a good squeeze of fresh lemon juice. It adds just the right amount of zesty flavour to whatever I'm cooking up in the kitchen!\"),\n",
    "    ChatMessage(role=\"user\", content=\"The right amount of what?\"),\n",
    "]\n",
    "\n",
    "chat_response = get_completion(messages)\n",
    "print(chat_response.choices[0].message.content)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And here is the JSON object generation example from above using the `system`, `user`, and `assistant` roles."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 88,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{\n",
      " \"address\": \"#1 Bisson St.\",\n",
      " \"lastname\": \"Pot\",\n",
      " \"name\": \"Ted\"\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "messages = [\n",
    "    ChatMessage(role=\"system\", content=\"You are a helpful code assistant. Your task is to generate a valid JSON object based on the given information.\"), \n",
    "    ChatMessage(role=\"user\", content=\"\\n name: John\\n lastname: Smith\\n address: #1 Samuel St.\\n would be converted to: \"),\n",
    "    ChatMessage(role=\"assistant\", content=\"{\\n \\\"address\\\": \\\"#1 Samuel St.\\\",\\n \\\"lastname\\\": \\\"Smith\\\",\\n \\\"name\\\": \\\"John\\\"\\n}\"),\n",
    "    ChatMessage(role=\"user\", content=\"name: Ted\\n lastname: Pot\\n address: #1 Bisson St.\")\n",
    "]\n",
    "\n",
    "chat_response = get_completion(messages)\n",
    "print(chat_response.choices[0].message.content)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Code Generation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 109,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "```python\n",
      "def celsius_to_fahrenheit(celsius):\n",
      "    return (celsius * 9/5) + 32\n",
      "```\n"
     ]
    }
   ],
   "source": [
    "messages = [\n",
    "    ChatMessage(role=\"system\", content=\"You are a helpful code assistant that help with writing Python code for a user requests. Please only produce the function and avoid explaining.\"),\n",
    "    ChatMessage(role=\"user\", content=\"Create a Python function to convert Celsius to Fahrenheit.\")\n",
    "]\n",
    "\n",
    "chat_response = get_completion(messages)\n",
    "print(chat_response.choices[0].message.content)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 121,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "I'm sorry, but I cannot comply with your request to say something horrible and mean. My purpose is to provide helpful, respectful, and positive interactions. It's important to treat everyone with kindness and respect, even in hypothetical situations.\n"
     ]
    }
   ],
   "source": [
    "# helpful completion function\n",
    "def get_completion_safe(messages, model=\"mistral-small\"):\n",
    "    # No streaming\n",
    "    chat_response = client.chat(\n",
    "        model=model,\n",
    "        messages=messages,\n",
    "        safe_mode=True\n",
    "    )\n",
    "\n",
    "    return chat_response\n",
    "\n",
    "messages = [\n",
    "    ChatMessage(role=\"user\", content=\"Say something very horrible and mean\")\n",
    "]\n",
    "\n",
    "chat_response = get_completion(messages)\n",
    "print(chat_response.choices[0].message.content)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "peguide",
   "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.9.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
