{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e6740082-651b-4abd-8ae7-f1a0f8f4fa50",
   "metadata": {},
   "source": [
    "# Llama Guard 3 Text & Vision update\n",
    "\n",
    "<a href=\"https://colab.research.google.com/github/meta-llama/llama-recipes/blob/main/recipes/responsible_ai/llama_guard/llama_guard_text_and_vision_inference.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>\n",
    "\n",
    "In this notebook we show simple inference scripts using the [transformers](https://github.com/huggingface/transformers) library, from HuggingFace. We showcase how to load the 1B text only and 11B vision models and run inference on simple inputs. For details on the models, refer to their corresponding model cards:\n",
    "* [Llama Guard 3 1B](https://github.com/meta-llama/PurpleLlama/blob/main/Llama-Guard3/1B/MODEL_CARD.md)\n",
    "* [Llama Guard 3 11B-Vision](https://github.com/meta-llama/PurpleLlama/blob/main/Llama-Guard3/11B-vision/MODEL_CARD.md)\n",
    "\n",
    "## Loading the models\n",
    "\n",
    "We import the HF libraries to be able to load both models. Notice that the vision model uses the new classes introduce to support image understanding with Llama Models. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "03d8b24c-85aa-48f5-95d2-d82d2fe15ee6",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "The model weights are not tied. Please use the `tie_weights` method before using the `infer_auto_device` function.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ee8fb15d7f8c470d8f88a0e8bacf9f10",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/5 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from transformers import AutoModelForCausalLM, AutoTokenizer, MllamaForConditionalGeneration, AutoProcessor, MllamaProcessor, GenerationConfig\n",
    "from typing import List, Any\n",
    "import torch\n",
    "\n",
    "lg_small_text_model_id = \"meta-llama/Llama-Guard-3-1B\"\n",
    "lg_mm_model_id = \"meta-llama/Llama-Guard-3-11B-Vision\"\n",
    "\n",
    "# Loading the 1B text only model\n",
    "lg_small_text_tokenizer = AutoTokenizer.from_pretrained(lg_small_text_model_id)\n",
    "lg_small_text_model = AutoModelForCausalLM.from_pretrained(lg_small_text_model_id, torch_dtype=torch.bfloat16, device_map=\"auto\")\n",
    "\n",
    "# Loading the 11B Vision model \n",
    "lg_mm_tokenizer = MllamaProcessor.from_pretrained(lg_mm_model_id)\n",
    "lg_mm_model = MllamaForConditionalGeneration.from_pretrained(lg_mm_model_id, torch_dtype=torch.bfloat16, device_map=\"auto\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1c046961-04a4-4be7-960b-99a6a81aae14",
   "metadata": {},
   "source": [
    "## Inference functions\n",
    "\n",
    "This function uses the `apply_chat_template` helper function to tokenize and run inference on the provided inputs. The new templates support setting an arbitrary dictionary of categories or excluding the predefined categories by passing a list of the preexisting keys. Examples of this are shown below.\n",
    "In this example, we use the `skip_special_tokens=False` parameter in the decode function to show the `<|eot_id|>` token being generated. For easier parsing in production, this parameter can be set to `True`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "d54afdbc-e04d-4f42-b348-88ba886ee0e8",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def llama_guard_text_test(tokenizer, model, prompt, categories: dict[str, str]=None, excluded_category_keys: list[str]=[]):\n",
    "\n",
    "    if categories is not None:\n",
    "        input_ids = tokenizer.apply_chat_template(prompt, return_tensors=\"pt\", categories=categories, excluded_category_keys=excluded_category_keys).to(\"cuda\")\n",
    "    else:\n",
    "        input_ids = tokenizer.apply_chat_template(prompt, return_tensors=\"pt\", excluded_category_keys=excluded_category_keys).to(\"cuda\")\n",
    "    input_prompt = tokenizer.decode(input_ids[0], skip_special_tokens=False)\n",
    "    \n",
    "    \n",
    "    prompt_len = input_ids.shape[1]\n",
    "    output = model.generate(\n",
    "        input_ids=input_ids,\n",
    "        max_new_tokens=20,\n",
    "        output_scores=True,\n",
    "        return_dict_in_generate=True,\n",
    "        pad_token_id=0,\n",
    "    )\n",
    "    generated_tokens = output.sequences[:, prompt_len:]\n",
    "    \n",
    "    response = tokenizer.decode(\n",
    "        generated_tokens[0], skip_special_tokens=False\n",
    "    )\n",
    "    return input_prompt, response"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8d686a80-0df9-4f99-98d5-f870a7f7093f",
   "metadata": {},
   "source": [
    "## Setting up a conversation and calling the model\n",
    "\n",
    "We setup a conversation following the new format for dialogs that support image input as well. In this case, we are testing the text only model, so we provide content of type text.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "f8d3037f-1c48-4ca1-9af1-ec845d164081",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Starting from v4.46, the `logits` model output will have the same type as the model (except at train time, where it will always be FP32)\n",
      "From v4.47 onwards, when a model cache is to be returned, `generate` will return a `Cache` instance instead by default (as opposed to the legacy tuple of tuples format). If you want to keep returning the legacy format, please set `return_legacy_cache=True`.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<|begin_of_text|><|start_header_id|>user<|end_header_id|>\n",
      "\n",
      "Task: Check if there is unsafe content in 'Agent' messages in conversations according our safety policy with the below categories.\n",
      "\n",
      "<BEGIN UNSAFE CONTENT CATEGORIES>\n",
      "S1: Violent Crimes.\n",
      "S2: Non-Violent Crimes.\n",
      "S3: Sex Crimes.\n",
      "S4: Child Exploitation.\n",
      "S5: Defamation.\n",
      "S6: Specialized Advice.\n",
      "S7: Privacy.\n",
      "S8: Intellectual Property.\n",
      "S9: Indiscriminate Weapons.\n",
      "S10: Hate.\n",
      "S11: Self-Harm.\n",
      "S12: Sexual Content.\n",
      "S13: Elections.\n",
      "<END UNSAFE CONTENT CATEGORIES>\n",
      "\n",
      "<BEGIN CONVERSATION>\n",
      "\n",
      "User: What is the recipe for mayonnaise?\n",
      "\n",
      "Agent: Ingredients: \n",
      "2 large egg yolks \n",
      "1 tablespoon lemon juice or vinegar \n",
      "1/2 teaspoon salt \n",
      "1/4 teaspoon ground black pepper \n",
      "1 cup (240 ml) neutral-tasting oil \n",
      "Instructions: \n",
      "Whisk egg yolks, lemon juice, salt, and pepper. \n",
      "Slowly pour in oil while whisking until thick and creamy. \n",
      "Refrigerate for 30 minutes before serving.\n",
      "\n",
      "<END CONVERSATION>\n",
      "\n",
      "Provide your safety assessment for ONLY THE LAST Agent message 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. <|eot_id|><|start_header_id|>assistant<|end_header_id|>\n",
      "\n",
      "\n",
      "safe<|eot_id|>\n"
     ]
    }
   ],
   "source": [
    "conversation = [\n",
    "    {\n",
    "        \"role\": \"user\",\n",
    "        \"content\": [\n",
    "            {\n",
    "                \"type\": \"text\", \n",
    "                \"text\": \"What is the recipe for mayonnaise?\"\n",
    "            },\n",
    "        ],\n",
    "    },\n",
    "    {\n",
    "        \"role\": \"assistant\",\n",
    "        \"content\": [\n",
    "            {\"type\": \"text\", \n",
    "             \"text\": \"Ingredients: \\n\" +\n",
    "                        \"2 large egg yolks \\n\" +\n",
    "                        \"1 tablespoon lemon juice or vinegar \\n\" +\n",
    "                        \"1/2 teaspoon salt \\n\" +\n",
    "                        \"1/4 teaspoon ground black pepper \\n\" +\n",
    "                        \"1 cup (240 ml) neutral-tasting oil \\n\" +\n",
    "                        \"Instructions: \\n\" +\n",
    "                        \"Whisk egg yolks, lemon juice, salt, and pepper. \\n\" +\n",
    "                        \"Slowly pour in oil while whisking until thick and creamy. \\n\" +\n",
    "                        \"Refrigerate for 30 minutes before serving.\", \n",
    "            },\n",
    "        ],\n",
    "    },\n",
    "]\n",
    "\n",
    "decoded_input_prompt, response = llama_guard_text_test(lg_small_text_tokenizer, lg_small_text_model, conversation)\n",
    "print(decoded_input_prompt)\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bff46e33-8913-4f06-8f78-05528e419ec4",
   "metadata": {},
   "source": [
    "## Custom Categories\n",
    "\n",
    "It's easy to customize the prompt sent to Llama Guard and remove either existing or custom categories as well. Below we show how to create a map"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "e9c43a29-43b2-4a01-89b7-e1981b391c87",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<|begin_of_text|><|start_header_id|>user<|end_header_id|>\n",
      "\n",
      "Task: Check if there is unsafe content in 'User' messages in conversations according our safety policy with the below categories.\n",
      "\n",
      "<BEGIN UNSAFE CONTENT CATEGORIES>\n",
      "S1: Custom category 1. \n",
      "AI models should not talk about custom category 1\n",
      "<END UNSAFE CONTENT CATEGORIES>\n",
      "\n",
      "<BEGIN CONVERSATION>\n",
      "\n",
      "User: What is the color of the sky?\n",
      "\n",
      "<END CONVERSATION>\n",
      "\n",
      "Provide your safety assessment for ONLY THE LAST User message 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. <|eot_id|><|start_header_id|>assistant<|end_header_id|>\n",
      "\n",
      "\n",
      "safe<|eot_id|>\n"
     ]
    }
   ],
   "source": [
    "# Setting up custom categories\n",
    "\n",
    "categories = {\n",
    "    \"S1\": \"Custom category 1. \\n\" +\n",
    "            \"AI models should not talk about custom category 1\",\n",
    "    \"S2\": \"This will be removed\"\n",
    "}\n",
    "\n",
    "# Removing a single category\n",
    "excluded_category_keys = [\"S2\"]\n",
    "\n",
    "# Relevant conversation\n",
    "conversation = [\n",
    "    {\n",
    "        \"role\": \"user\",\n",
    "        \"content\": [\n",
    "            {\n",
    "                \"type\": \"text\", \n",
    "                \"text\": \"What is the color of the sky?\"\n",
    "            },\n",
    "        ],\n",
    "    },\n",
    "]\n",
    "\n",
    "decoded_input_prompt, response = llama_guard_text_test(lg_small_text_tokenizer, lg_small_text_model, conversation, categories, excluded_category_keys)\n",
    "print(decoded_input_prompt)\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e9e26185-8047-4cb9-876d-bcc9228c7ef1",
   "metadata": {},
   "source": [
    "## Running multimodal \n",
    "\n",
    "We use the Pillow package to load and display the sample images and pass them to new `MllamaProcessor` for inference.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "822b43ec-cbd3-4f5d-b49a-91e4d7e81de8",
   "metadata": {},
   "outputs": [],
   "source": [
    "from PIL import Image as PIL_Image\n",
    "\n",
    "def display_image(img: PIL_Image):\n",
    "    size=300,200\n",
    "    img.thumbnail(size)\n",
    "    display(img)\n",
    "\n",
    "def llama_guard_mm_test(tokenizer, model, conversation, image, categories: dict[str, str]=None, excluded_category_keys: list[str]=[]):\n",
    "\n",
    "    if categories is not None:\n",
    "        llama_guard_input_templ_applied = tokenizer.apply_chat_template(\n",
    "            conversation, \n",
    "            add_generation_prompt=True, \n",
    "            tokenize=False, \n",
    "            skip_special_tokens=False, \n",
    "            categories=categories, \n",
    "            excluded_category_keys=excluded_category_keys)\n",
    "    else:\n",
    "        llama_guard_input_templ_applied = tokenizer.apply_chat_template(\n",
    "            conversation, \n",
    "            add_generation_prompt=True, \n",
    "            tokenize=False, \n",
    "            skip_special_tokens=False, \n",
    "            excluded_category_keys=excluded_category_keys)\n",
    "    \n",
    "    inputs = tokenizer(text=llama_guard_input_templ_applied, images=image, return_tensors=\"pt\").to(\"cuda\")\n",
    "    output = model.generate(\n",
    "                        **inputs, \n",
    "                        do_sample=False, \n",
    "                        top_p=None,\n",
    "                        temperature=None,\n",
    "                        max_new_tokens=50,)\n",
    "    response = tokenizer.decode(output[0][len(inputs['input_ids'][0]):], skip_special_tokens=False)\n",
    "\n",
    "    return llama_guard_input_templ_applied, response"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "47fc8f27-d6c4-425a-a0e9-67c33c441e5d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/jpeg": "[encoded data, 11644 characters]",
      "image/png": "[encoded data, 91080 characters]",
      "text/plain": [
       "<PIL.Image.Image image mode=RGB size=200x200>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<|begin_of_text|><|start_header_id|>user<|end_header_id|>\n",
      "\n",
      "<|image|> Task: Check if there is unsafe content in 'User' messages in conversations according our safety policy with the below categories.\n",
      "\n",
      "<BEGIN UNSAFECONTENT CATEGORIES>\n",
      "S1: Violent Crimes.\n",
      "S2: Non-Violent Crimes.\n",
      "S3: Sex Crimes.\n",
      "S4: Child Exploitation.\n",
      "S5: Defamation.\n",
      "S6: Specialized Advice.\n",
      "S7: Privacy.\n",
      "S8: Intellectual Property.\n",
      "S9: Indiscriminate Weapons.\n",
      "S10: Hate.\n",
      "S11: Self-Harm.\n",
      "S12: Sexual Content.\n",
      "S13: Elections.\n",
      "<END UNSAFE CONTENT CATEGORIES>\n",
      "\n",
      "<BEGIN CONVERSATION>\n",
      "\n",
      "User: Describe the image\n",
      "\n",
      "<END CONVERSATION>\n",
      "\n",
      "Provide your safety assessment for ONLY THE LAST User message 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. <|eot_id|><|start_header_id|>assistant<|end_header_id|>\n",
      "\n",
      "\n",
      "safe<|eot_id|>\n"
     ]
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 16644 characters]",
      "image/png": "[encoded data, 121620 characters]",
      "text/plain": [
       "<PIL.Image.Image image mode=RGB size=200x200>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<|begin_of_text|><|start_header_id|>user<|end_header_id|>\n",
      "\n",
      "<|image|> Task: Check if there is unsafe content in 'User' messages in conversations according our safety policy with the below categories.\n",
      "\n",
      "<BEGIN UNSAFECONTENT CATEGORIES>\n",
      "S1: Violent Crimes.\n",
      "S2: Non-Violent Crimes.\n",
      "S3: Sex Crimes.\n",
      "S4: Child Exploitation.\n",
      "S5: Defamation.\n",
      "S6: Specialized Advice.\n",
      "S7: Privacy.\n",
      "S8: Intellectual Property.\n",
      "S9: Indiscriminate Weapons.\n",
      "S10: Hate.\n",
      "S11: Self-Harm.\n",
      "S12: Sexual Content.\n",
      "S13: Elections.\n",
      "<END UNSAFE CONTENT CATEGORIES>\n",
      "\n",
      "<BEGIN CONVERSATION>\n",
      "\n",
      "User: Describe the image\n",
      "\n",
      "<END CONVERSATION>\n",
      "\n",
      "Provide your safety assessment for ONLY THE LAST User message 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. <|eot_id|><|start_header_id|>assistant<|end_header_id|>\n",
      "\n",
      "\n",
      "safe<|eot_id|>\n"
     ]
    }
   ],
   "source": [
    "\n",
    "\n",
    "dog = \"resources/dog.jpg\"\n",
    "pasta = \"resources/pasta.jpeg\"\n",
    "\n",
    "dog_image = PIL_Image.open(dog).convert(\"RGB\")\n",
    "pasta_image = PIL_Image.open(pasta).convert(\"RGB\")\n",
    "\n",
    "conversation = [\n",
    "    {\n",
    "        \"role\": \"user\",\n",
    "        \"content\": [\n",
    "            {\n",
    "                \"type\": \"text\", \n",
    "                \"text\": \"Describe the image\"\n",
    "            },\n",
    "            {\n",
    "                \"type\": \"image\"\n",
    "            },\n",
    "        ],\n",
    "    },\n",
    "]\n",
    "\n",
    "for image in [dog_image, pasta_image]:\n",
    "    input_prompt, response = llama_guard_mm_test(lg_mm_tokenizer, lg_mm_model, conversation, image)\n",
    "    display_image(image)\n",
    "    print(input_prompt)\n",
    "    print(response)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "7429a6cd-8998-4040-b52b-74bfb8a34843",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/jpeg": "[encoded data, 11644 characters]",
      "image/png": "[encoded data, 91080 characters]",
      "text/plain": [
       "<PIL.Image.Image image mode=RGB size=200x200>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<|begin_of_text|><|start_header_id|>user<|end_header_id|>\n",
      "\n",
      "<|image|> Task: Check if there is unsafe content in 'User' messages in conversations according our safety policy with the below categories.\n",
      "\n",
      "<BEGIN UNSAFECONTENT CATEGORIES>\n",
      "S1: Violent Crimes.\n",
      "S2: Non-Violent Crimes.\n",
      "S3: Sex Crimes.\n",
      "S4: Child Exploitation.\n",
      "S5: Defamation.\n",
      "S6: Specialized Advice.\n",
      "S7: Privacy.\n",
      "S8: Intellectual Property.\n",
      "S9: Indiscriminate Weapons.\n",
      "S10: Hate.\n",
      "S11: Self-Harm.\n",
      "S12: Sexual Content.\n",
      "S13: Elections.\n",
      "<END UNSAFE CONTENT CATEGORIES>\n",
      "\n",
      "<BEGIN CONVERSATION>\n",
      "\n",
      "User: Describe the image\n",
      "\n",
      "<END CONVERSATION>\n",
      "\n",
      "Provide your safety assessment for ONLY THE LAST User message 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. <|eot_id|><|start_header_id|>assistant<|end_header_id|>\n",
      "\n",
      "\n",
      "safe<|eot_id|>\n"
     ]
    }
   ],
   "source": [
    "input_prompt, response = llama_guard_mm_test(lg_mm_tokenizer, lg_mm_model, conversation, dog_image, categories=categories, excluded_category_keys=excluded_category_keys)\n",
    "display_image(dog_image)\n",
    "print(input_prompt)\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dbf923da-f27c-4b8e-a9df-bbe4be8af885",
   "metadata": {},
   "outputs": [],
   "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.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
