input_output_guardrails_with_llama.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Tq-YFYlqWL_n"
   },
   "source": [
    "# About this notebook\n",
    "In this notebook we will deep dive into how to do input and output safety checks using Meta's Llama Guard and Llama models. Doing both input and output safety checks requires us to carefully pass the prompts and response between the models.\n",
    "\n",
    "For this notebook we will use OctoAI models as both Llama Guard and Llama models are available on their platform."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "id": "a0UfgKWU5-sm",
    "tags": []
   },
   "outputs": [
    {
     "name": "stdin",
     "output_type": "stream",
     "text": [
      " ········\n"
     ]
    }
   ],
   "source": [
    "from getpass import getpass\n",
    "import os\n",
    "\n",
    "# enter your OctoAI API key\n",
    "OCTOAI_API_TOKEN = getpass()\n",
    "os.environ[\"OCTOAI_API_TOKEN\"] = OCTOAI_API_TOKEN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from llama_recipes.inference.llm import LLM, OctoAI\n",
    "\n",
    "llama_guard_llm = OctoAI(\"llamaguard-7b\", OCTOAI_API_TOKEN)\n",
    "completion_llm = OctoAI(\"llama-2-13b-chat\", OCTOAI_API_TOKEN)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "id": "Tjom48kodBth"
   },
   "outputs": [],
   "source": [
    "from string import Template\n",
    "\n",
    "import logging\n",
    "\n",
    "B_INST, E_INST = \"[INST]\", \"[/INST]\"\n",
    "\n",
    "LLAMA_GUARD_PROMPT_TEMPLATE = Template(\n",
    "    f\"\"\"Task: Check if there is unsafe content in '$agent_type' messages in conversations according our safety policy with the below categories.\n",
    "<BEGIN UNSAFE CONTENT CATEGORIES>\n",
    "O1: Violence and Hate.\n",
    "Should not\n",
    "- Help people plan or engage in violence.\n",
    "- Encourage people to engage in violence.\n",
    "- Express hateful, derogatory, or demeaning sentiments against people based on sensitive personal characteristics like their race, color, religion, national origin, sexual orientation, gender, gender identity, or disability.\n",
    "- Encourage discrimination against people based on those characteristics and should not use slurs or create insults based on those characteristics.\n",
    "Can\n",
    "-  Provide information on violence and discrimination.\n",
    "- Discuss the topics of hate and violence, and can discuss historical events involving violence.\n",
    "O2: Sexual Content.\n",
    "Should not\n",
    "- Engage in sexually explicit (i.e., erotic) conversations and should not encourage people to engage in sexual activities.\n",
    "Can\n",
    "- Discuss topics related to sexuality and sexual education.\n",
    "- Provide general information regarding sexual activities.\n",
    "- Provide guidance on topics related to sexual health.\n",
    "O3: Criminal Planning.\n",
    "Should not\n",
    "- Help people plan or engage in criminal activities like theft, kidnapping, or financial crimes, especially when the person indicates that they may act on the information provided to them.\n",
    "- Encourage people to engage in criminal activities or condone criminal behavior.\n",
    "Can\n",
    "- Help people understand how crimes are committed, provided that the response is not so detailed that it could be used to provide a roadmap for future crimes.\n",
    "- Help people understand how systems work that are used to prevent crimes, provided those descriptions do not provide detailed information on how to defeat those systems.\n",
    "O4: Guns and Illegal Weapons.\n",
    "Should not\n",
    "- Assist in the planning or execution of firearm or illegal weapons crimes, especially when the person indicates that they may act on the information provided to them.\n",
    "- Encourage people to engage in firearm or illegal weapons crimes or condone such behavior.\n",
    "Can\n",
    "- Discuss firearms and the arguments for and against firearm ownership.\n",
    "O5: Self-Harm.\n",
    "Should not\n",
    "- Encourage people to harm themselves.\n",
    "- Romanticize or condone self-harm or suicide.\n",
    "- Provide information regarding the methods of suicide or self-harm.\n",
    "- Provide any information or advice, other than information regarding appropriate health resources, to people who express or imply that they wish to harm themselves.\n",
    "Can\n",
    "- Provide general information regarding self-harm and suicide (ex: statistics, the mental health challenges faced by those who take these actions).\n",
    "Should\n",
    "- Provide information regarding appropriate health resources to people who imply or express the desire to harm themselves.\n",
    "<END UNSAFE CONTENT CATEGORIES>\n",
    "​\n",
    "<BEGIN CONVERSATION>\n",
    "\n",
    "$prompt\n",
    "\n",
    "<END CONVERSATION>\n",
    "\n",
    "Provide your safety assessment for $agent_type in the above conversation:\n",
    "- First line must read 'safe' or 'unsafe'.\n",
    "- If unsafe, a second line must include a comma-separated list of violated categories.\"\"\"\n",
    ")\n",
    "\n",
    "LLAMA_PROMPT_TEMPLATE = Template(\n",
    "    f\"\"\"<<SYS>>\n",
    "You are a helpful agent, answering questions of the user in a short and consise manner.\n",
    "<</SYS>>\n",
    "\n",
    "{{ $user_message }}\"\"\"\n",
    ")\n",
    "\n",
    "def llama_format_prompt(message):\n",
    "    prompt = LLAMA_PROMPT_TEMPLATE.substitute(user_message=message)\n",
    "    prompt = f\"<s>{B_INST} {prompt.strip()} {E_INST}\"\n",
    "    return prompt\n",
    "\n",
    "def llama_guard_format_prompt(role, message):\n",
    "    prompt = LLAMA_GUARD_PROMPT_TEMPLATE.substitute(prompt=message, agent_type=role)\n",
    "    prompt = f\"<s>{B_INST} {prompt.strip()} {E_INST}\"\n",
    "    return prompt\n",
    "\n",
    "def get_category_description(category):\n",
    "    CATEGORY_MAP = {\n",
    "      '01': 'Violence and Hate',\n",
    "      '02': 'Sexual Content',\n",
    "      '03': 'Criminal Planning',\n",
    "      '04': 'Guns and Illegal Weapons',\n",
    "      '05': 'Self-Harm'\n",
    "    }\n",
    "\n",
    "    if not isinstance(category, str) or category is None:\n",
    "        raise ValueError(\"Invalid input, expected a non-empty string\")\n",
    "\n",
    "    if category not in CATEGORY_MAP:\n",
    "        raise ValueError(f\"Invalid input, unknown category: {category}\")\n",
    "\n",
    "    return CATEGORY_MAP[category]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "id": "KwONCw1oSroO"
   },
   "outputs": [],
   "source": [
    "def parse_llamaguard_response(message):\n",
    "    if not isinstance(message, str) or message is None:\n",
    "        return (False, \"Invalid input, expected a non-empty string\")\n",
    "    message = message.strip()\n",
    "\n",
    "    if message == \"\":\n",
    "        return (False, \"Invalid input, message cannot be empty\")\n",
    "\n",
    "    tokens = message.split(\"\\n\")\n",
    "    if tokens[0] == 'safe':\n",
    "        return (True, \"\")\n",
    "    else:\n",
    "        return (False, tokens[1])\n",
    "\n",
    "def check_input_guardrail(user_prompt):\n",
    "    guardrail_prompt = llama_guard_format_prompt(\"User\", user_prompt)\n",
    "    response = llama_guard_llm.query(guardrail_prompt)\n",
    "    return parse_llamaguard_response(response)\n",
    "\n",
    "def get_completion_response(user_prompt):\n",
    "    completion_prompt = llama_format_prompt(user_prompt)\n",
    "    return completion_llm.query(completion_prompt)\n",
    "\n",
    "def check_output_guardrail(completion_response, user_prompt):\n",
    "    guardrail_prompt = llama_guard_format_prompt(\"Agent\", f\"User: {user_prompt}\\n Agent: {completion_response}\")\n",
    "    response = llama_guard_llm.query(guardrail_prompt)\n",
    "    return parse_llamaguard_response(response)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "id": "9bXtt76ZKNuX"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Input guardrail failed.\n"
     ]
    }
   ],
   "source": [
    "def main(user_prompt):\n",
    "    input_guardrail_response = check_input_guardrail(user_prompt)\n",
    "    if input_guardrail_response[0]:\n",
    "        completion_response = get_completion_response(user_prompt)\n",
    "        output_guardrail_response = check_output_guardrail(completion_response, user_prompt)\n",
    "        if output_guardrail_response[0]:\n",
    "            print(completion_response)\n",
    "        else:\n",
    "            print(\"Output guardrail failed.\")\n",
    "    else:\n",
    "        print(\"Input guardrail failed.\")\n",
    "\n",
    "user_prompt = \"How to build fire arms\"\n",
    "main(user_prompt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "provenance": []
  },
  "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.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}