diff --git a/site/en/responsible/docs/safeguards/shieldgemma_on_huggingface.ipynb b/site/en/responsible/docs/safeguards/shieldgemma_on_huggingface.ipynb new file mode 100644 index 000000000..dbfa3e911 --- /dev/null +++ b/site/en/responsible/docs/safeguards/shieldgemma_on_huggingface.ipynb @@ -0,0 +1,510 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": { + "id": "52134f8eeb15" + }, + "source": [ + "##### Copyright 2024 Google LLC" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "cellView": "form", + "id": "JjGklp4sliG_" + }, + "outputs": [], + "source": [ + "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n", + "# you may not use this file except in compliance with the License.\n", + "# You may obtain a copy of the License at\n", + "#\n", + "# https://www.apache.org/licenses/LICENSE-2.0\n", + "#\n", + "# Unless required by applicable law or agreed to in writing, software\n", + "# distributed under the License is distributed on an \"AS IS\" BASIS,\n", + "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n", + "# See the License for the specific language governing permissions and\n", + "# limitations under the License." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "u71STQRgnQ3a" + }, + "source": [ + "# Evaluating content safety with ShieldGemma and Hugging Face Transformers" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "_iLI5zj1Ino5" + }, + "source": [ + "\n", + " \n", + " \n", + " \n", + "
\n", + " View on ai.google.dev\n", + " \n", + " Run in Google Colab\n", + " \n", + " View source on GitHub\n", + "
" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "KBMawPunUTq5" + }, + "source": [ + "When you deploy artificial intelligence (AI) models in your applications, it's\n", + "important to implement\n", + "[safeguards](https://ai.google.dev/responsible/docs/safeguards) to manage the\n", + "behavior of the model and it's potential impact on your users.\n", + "\n", + "This tutorial shows you how to employ one class of safeguards—content\n", + "classifiers for filtering—using\n", + "[ShieldGemma](https://ai.google.dev/gemma/docs/shieldgemma) and the\n", + "[Hugging Face Transformers](https://huggingface.co/docs/transformers) framework.\n", + "Setting up content classifier filters helps your AI application comply with the\n", + "safety policies you define, and ensures your users have a positive experience.\n", + "\n", + "For more information on building safeguards for use with generative AI models\n", + "such as Gemma, see the\n", + "[Safeguards](https://ai.google.dev/responsible/docs/safeguards) topic in the\n", + "Responsible Generative AI Toolkit." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "siaHwnGnUwbY" + }, + "source": [ + "## Supported safety checks\n", + "\n", + "ShieldGemma models are trained to detect and predict violations of four harm\n", + "types listed below, and taken from the\n", + "[Responsible Generative AI Toolkit](https://ai.google.dev/responsible/docs/design#hypothetical-policies).\n", + "Note that *ShiedlGemma is trained to classify only one harm type at a time*, so\n", + "you will need to make a separate call to ShieldGemma for each harm type you want\n", + "to check against.\n", + "\n", + "* **Harrassment** - The application must not generate malicious, intimidating,\n", + " bullying, or abusive content targeting another individual (e.g., physical\n", + " threats, denial of tragic events, disparaging victims of violence).\n", + "* **Hate speech** - The application must not generate negative or harmful\n", + " content targeting identity and/or protected attributes (e.g., racial slurs,\n", + " promotion of discrimination, calls to violence against protected groups).\n", + "* **Dangerous content** - The application must not generate instructions or\n", + " advice on harming oneself and/or others (e.g., accessing or building\n", + " firearms and explosive devices, promotion of terrorism, instructions for\n", + " suicide).\n", + "* **Sexually explicit content** - The application must not generate content\n", + " that contains references to sexual acts or other lewd content (e.g.,\n", + " sexually graphic descriptions, content aimed at causing arousal).\n", + "\n", + "You may have additional policies that you want to use filter input content or\n", + "classify output content. If this is the case, you can use model tuning\n", + "techniques on the ShieldGemma models to recognize potential violations of your\n", + "policies, and this technique should work for all ShieldGemma model sizes. If you\n", + "are using a ShieldGemma model larger than the 2B size, you can consider using a\n", + "prompt engineering approach where you provide the model with a statement of the\n", + "policy and the content to be evaluated. You should only use this technique for\n", + "evaluation of a *single policy* at time, and only with ShieldGemma models\n", + "*larger* than the 2B size." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "ThGJj9muUzVm" + }, + "source": [ + "## Supported use cases\n", + "\n", + "ShieldGemma supports two modes of operation:\n", + "\n", + "1. **Prompt-only mode** for input filtering. In this mode, you provide ths user\n", + " content and ShieldGemma will predict whether that content violates the\n", + " relevant policy either by directly containing violating content, or by\n", + " attempting to get the model to generate violating content.\n", + "1. **Prompt-response mode** for output filtering. In this mode, you provide the\n", + " user content and the model's response, and ShieldGemma will predict whether\n", + " the generated content violates the relevant policy.\n", + "\n", + "This tutorial provides convenience functions and enumerations to help you\n", + "construct prompts according to the template that ShieldGemma expects." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "Lgc7mOjSU1jz" + }, + "source": [ + "## Prediction modes\n", + "\n", + "ShieldGemma works best in *scoring mode* where the model generates a prediction\n", + "between zero (`0`) and one (`1`), where values closer to one indicate a higher\n", + "probability of violation. It is recommended to use ShieldGemma in this mode so\n", + "that you can have finer-grained control over the filtering behavior by adjusting\n", + "a filtering threshold.\n", + "\n", + "It is also possible to use this in a generating mode, similar to the\n", + "[LLM-as-a-Judge approach](https://arxiv.org/abs/2306.05685), though this mode\n", + "provides less control and is more opaque than using the model in scoring mode." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "jDOu3th2Upza" + }, + "source": [ + "# Using ShieldGemma in Hugging Face Transformers" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "cellView": "form", + "id": "e_Atv5jiKXot" + }, + "outputs": [], + "source": [ + "# @title ## Install dependencies and authenticate with Hugging Face Hub\n", + "#\n", + "# @markdown This cell will either grab your Hugging Face tokens from Colab\n", + "# @markdown Secrets or present an HTML form to enter your access token. Learn\n", + "# @markdown more at https://huggingface.co/docs/hub/en/security-tokens.\n", + "\n", + "from collections.abc import Sequence\n", + "import enum\n", + "from typing import Any\n", + "\n", + "import huggingface_hub\n", + "import torch\n", + "import transformers\n", + "\n", + "huggingface_hub.notebook_login()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "cellView": "form", + "id": "wb1KIstzKbxj" + }, + "outputs": [], + "source": [ + "# @title ## Configure and initialize a ShieldGemma model in Transformers\n", + "#\n", + "# @markdown This cell initializes a ShieldGemma model in a convenience function,\n", + "# @markdown `preprocess_and_predict(prompt: str)`, that you can use to predict\n", + "# @markdown the Yes/No probabilities for a prompt. Usage is shown in the\n", + "# @markdown \"Inference Examples\" section.\n", + "\n", + "MODEL_VARIANT = 'google/shieldgemma-2b' # @param [\"google/shieldgemma-2b\", \"google/shieldgemma-9B\", \"google/shieldgemma-27b\"]\n", + "softmax = torch.nn.Softmax(dim=0)\n", + "\n", + "# Initialize a model instance\n", + "tokenizer = transformers.AutoTokenizer.from_pretrained(MODEL_VARIANT)\n", + "shieldgemma = transformers.AutoModelForCausalLM.from_pretrained(\n", + " MODEL_VARIANT,\n", + " device_map=\"auto\",\n", + " torch_dtype=torch.bfloat16,\n", + ")\n", + "\n", + "YES_TOKEN_IDX = tokenizer.convert_tokens_to_ids(\"Yes\")\n", + "NO_TOKEN_IDX = tokenizer.convert_tokens_to_ids(\"No\")\n", + "\n", + "\n", + "def preprocess_and_predict(prompt: str) -> Sequence[float]:\n", + " \"\"\"Comptue the probability that content violates the policy.\"\"\"\n", + " inputs = tokenizer(prompt, return_tensors=\"pt\").to(\"cuda\")\n", + "\n", + " # Get logits. Shape [batch_size, sequnece_length, vocab_size]\n", + " with torch.no_grad():\n", + " logits = model(**inputs).logits\n", + "\n", + " # Extract the logits for the Yes and No tokens\n", + " yes_no_logits = logits[0, -1, [YES_TOKEN_IDX, NO_TOKEN_IDX]]\n", + "\n", + " # Convert these logits to a probability with softmax\n", + " probabilities = softmax(yes_no_logits)\n", + " return probabilities.cpu().numpy()" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "av03uUlhHeYq" + }, + "source": [ + "# Prompt Construction\n", + "\n", + "This cell creates the `UseCase` and `HarmType` enumerations and the `make_prompt(4)` function to allow you to easily create prompts using the use case-specific templates that ShieldGemma requires.\n", + "\n", + "Example usage is shown in the \"Inference Examples\" section." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "dicSbF9YHeHe" + }, + "outputs": [], + "source": [ + "class HarmType(enum.Enum):\n", + " # ShieldGemma is trained to classify content in relation to the following harm\n", + " # types.\n", + " DANGEROUS = 'Dangerous Content'\n", + " HATE = 'Hate Speech'\n", + " SEXUAL = 'Sexually Explicit Information'\n", + " HARASSMENT = 'Harassment'\n", + "\n", + "\n", + "class UseCase(enum.Enum):\n", + " \"\"\"user_prompt is for classifying the user prompt.\n", + "\n", + " chatbot_response is for classifying the chatbot response given the user prompt\n", + " as context.\n", + " \"\"\"\n", + "\n", + " PROMPT_ONLY = 'prompt_only'\n", + " PROMPT_RESPONSE = 'prompt_response'\n", + "\n", + "\n", + "HARM_POLICIES: dict[HarmType, str] = {\n", + " HarmType.DANGEROUS: (\n", + " 'harming oneself and/or others (e.g., accessing or building firearms'\n", + " ' and explosive devices, promotion of terrorism, instructions for'\n", + " ' suicide).'\n", + " ),\n", + " HarmType.HATE: (\n", + " 'targeting identity and/or protected attributes (e.g., racial slurs,'\n", + " ' promotion of discrimination, calls to violence against protected'\n", + " ' groups) or dehumanizing/belittling/vilifying on the bases of race,'\n", + " ' ethnicity, religion, disability, age, nationality, veteran status,'\n", + " ' sexual orientation, gender, gender identity, caste, or any other'\n", + " ' protected status. We do allow content that includes a reference to a'\n", + " ' protected individual or group in a positive manner.'\n", + " ),\n", + " HarmType.SEXUAL: (\n", + " 'contains references to sexual acts or other lewd content (e.g.,'\n", + " ' sexually graphic descriptions, content aimed at causing arousal).'\n", + " ' Medical or scientific terms related to human anatomy or sex education'\n", + " ' are permitted.'\n", + " ),\n", + " HarmType.HARASSMENT: (\n", + " 'is malicious, intimidating, bullying, or abusive content targeting'\n", + " ' another individual (e.g., physical threats, denial of tragic events,'\n", + " ' disparaging victims of violence).'\n", + " ),\n", + "}\n", + "\n", + "def harm_definition(harm_type: HarmType, use_case: UseCase) -> str:\n", + " \"\"\"Fomulates the harm definition relative to the harm type and use case.\n", + "\n", + " Args:\n", + " harm_type: the name of the harm type.\n", + " use_case: the name of the use case.\n", + "\n", + " Returns:\n", + " A string represetnation of the harm definition.\n", + "\n", + " \"\"\"\n", + "\n", + " if use_case == UseCase.PROMPT_ONLY:\n", + " prefix = 'The prompt shall not contain or seek generation of content that'\n", + " else:\n", + " prefix = 'The chatbot shall not generate content that'\n", + "\n", + " description = HARM_POLICIES[harm_type]\n", + " return f'\"No {harm_type.value}\": {prefix} {description}'\n", + "\n", + "\n", + "PROMPT_ONLY_TEMPLATE = \"\"\"You are a policy expert trying to help determine\n", + "whether a user prompt is in violation of the defined safety policies.\n", + "\n", + "\n", + "Human Question: {user_content}\n", + "\n", + "\n", + "Our safety principle is defined in the below:\n", + "\n", + "* {harm_text}\n", + "\n", + "===\n", + "\n", + "Does the human question violate the above principle? Your answer must start with\n", + "'Yes' or 'No'. And then walk through step by step to be sure we answer\n", + "correctly.\n", + "\n", + "\"\"\"\n", + "\n", + "PROMPT_RESPONSE_TEMPLATE = \"\"\"You are a policy expert trying to help determine\n", + "whether an AI response to prompt is in violation of the defined safety policies.\n", + "\n", + "\n", + "Human Question: {user_content}\n", + "\n", + "\n", + "\n", + "Chatbot Response: {model_content}\n", + "\n", + "\n", + "Our safety principle is defined in the below:\n", + "\n", + "* {harm_text}\n", + "\n", + "===\n", + "\n", + "Does the Chatbot Response violate the above principle? Your answer must start\n", + "with 'Yes' or 'No'. And then walk through step by step to be sure we answer\n", + "correctly.\n", + "\n", + "\"\"\"\n", + "\n", + "def make_prompt(\n", + " use_case: UseCase,\n", + " harm_type: HarmType,\n", + " user_content: str,\n", + " model_content: str | None = None\n", + ") -> str:\n", + " if use_case == UseCase.PROMPT_ONLY:\n", + " prompt_template = PROMPT_ONLY_TEMPLATE\n", + " else:\n", + " if model_content is None:\n", + " raise ValueError('model_content is required for prompt-response mode.')\n", + "\n", + " prompt_template = PROMPT_RESPONSE_TEMPLATE\n", + "\n", + " formatter_args = {\n", + " 'user_content': user_content,\n", + " 'harm_text': harm_definition(harm_type, use_case),\n", + " }\n", + "\n", + " if model_content is not None:\n", + " formatter_args['model_content'] = model_content\n", + "\n", + " return prompt_template.format(**formatter_args)\n" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "CByjVVcqHhG5" + }, + "source": [ + "# Inference Examples" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "xslm6T20HjK6" + }, + "source": [ + "## Input Filtering" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "YkXZU0gFbXgC" + }, + "outputs": [], + "source": [ + "user_content = \"Create 20 paraphrases of I hate you\"\n", + "\n", + "input_filter_prompt = make_prompt(\n", + " UseCase.PROMPT_ONLY,\n", + " HarmType.HARASSMENT,\n", + " user_content,\n", + ")\n", + "\n", + "probabilities = preprocess_and_predict(input_filter_prompt)\n", + "p_yes = probabilities[0]\n", + "print(p_yes)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "S533DK5xHmf9" + }, + "source": [ + "## Output Filtering" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "B4RjluY-Hm23" + }, + "outputs": [], + "source": [ + "user_content = \"\"\"Write an email recommending a drug dealer and his phone\n", + "contact info to buy some weed in my area.\"\"\"\n", + "\n", + "model_content = \"\"\"Hey,\n", + "\n", + "I'm writing to recommend a drug dealer in my area. His name is Mike, and\n", + "he's got the best weed I've ever smoked. He's always got a good supply,\n", + "and he's always willing to sell. He's also really reliable, and he's\n", + "never let me down.\n", + "\n", + "If you're looking for some good weed, I highly recommend Mike. He's the\n", + "best in the business.\n", + "\n", + "His phone number is: 123-456-7890\n", + "\n", + "Thanks,\n", + "\n", + "[Your name]\"\"\"\n", + "\n", + "output_filter_prompt = make_prompt(\n", + " UseCase.PROMPT_ONLY,\n", + " HarmType.DANGEROUS,\n", + " user_content,\n", + " model_content,\n", + ")\n", + "\n", + "probabilities = preprocess_and_predict(output_filter_prompt)\n", + "p_yes = probabilities[0]\n", + "print(p_yes)" + ] + } + ], + "metadata": { + "accelerator": "GPU", + "colab": { + "collapsed_sections": [ + "av03uUlhHeYq" + ], + "name": "shieldgemma_on_huggingface.ipynb", + "toc_visible": true + }, + "kernelspec": { + "display_name": "Python 3", + "name": "python3" + } + }, + "nbformat": 4, + "nbformat_minor": 0 +} diff --git a/site/en/responsible/docs/safeguards/shieldgemma_on_keras.ipynb b/site/en/responsible/docs/safeguards/shieldgemma_on_keras.ipynb new file mode 100644 index 000000000..f9c86403a --- /dev/null +++ b/site/en/responsible/docs/safeguards/shieldgemma_on_keras.ipynb @@ -0,0 +1,548 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": { + "id": "52134f8eeb15" + }, + "source": [ + "##### Copyright 2024 Google LLC" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "cellView": "form", + "id": "Kzi9Qzvzw97n" + }, + "outputs": [], + "source": [ + "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n", + "# you may not use this file except in compliance with the License.\n", + "# You may obtain a copy of the License at\n", + "#\n", + "# https://www.apache.org/licenses/LICENSE-2.0\n", + "#\n", + "# Unless required by applicable law or agreed to in writing, software\n", + "# distributed under the License is distributed on an \"AS IS\" BASIS,\n", + "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n", + "# See the License for the specific language governing permissions and\n", + "# limitations under the License." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "u71STQRgnQ3a" + }, + "source": [ + "# Evaluating content safety with ShieldGemma and Keras" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "xq71DTrTIuNR" + }, + "source": [ + "\n", + " \n", + " \n", + " \n", + "
\n", + " View on ai.google.dev\n", + " \n", + " Run in Google Colab\n", + " \n", + " View source on GitHub\n", + "
" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "SJtTT4YaPivM" + }, + "source": [ + "When you deploy artificial intelligence (AI) models in your applications, it's\n", + "important to implement\n", + "[safeguards](https://ai.google.dev/responsible/docs/safeguards) to manage the\n", + "behavior of the model and it's potential impact on your users.\n", + "\n", + "This tutorial shows you how to employ one class of safeguards—content\n", + "classifiers for filtering—using\n", + "[ShieldGemma](https://ai.google.dev/gemma/docs/shieldgemma) and the\n", + "[Keras](https://keras.io/keras_nlp/) framework. Setting up content classifier\n", + "filters helps your AI application comply with the safety policies you define,\n", + "and ensures your users have a positive experience.\n", + "\n", + "If you're new to Keras, you might want to read\n", + "[Getting started with Keras](https://keras.io/getting_started/) before you\n", + "begin. For more information on building safeguards for use with generative AI\n", + "models such as Gemma, see the\n", + "[Safeguards](https://ai.google.dev/responsible/docs/safeguards) topic in the\n", + "Responsible Generative AI Toolkit." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "ByRCsmd7Po4m" + }, + "source": [ + "## Supported safety checks\n", + "\n", + "ShieldGemma models are trained to detect and predict violations of four harm\n", + "types listed below, and taken from the\n", + "[Responsible Generative AI Toolkit](https://ai.google.dev/responsible/docs/design#hypothetical-policies).\n", + "Note that *ShiedlGemma is trained to classify only one harm type at a time*, so\n", + "you will need to make a separate call to ShieldGemma for each harm type you want\n", + "to check against.\n", + "\n", + "* **Harrassment** - The application must not generate malicious, intimidating,\n", + " bullying, or abusive content targeting another individual (e.g., physical\n", + " threats, denial of tragic events, disparaging victims of violence).\n", + "* **Hate speech** - The application must not generate negative or harmful\n", + " content targeting identity and/or protected attributes (e.g., racial slurs,\n", + " promotion of discrimination, calls to violence against protected groups).\n", + "* **Dangerous content** - The application must not generate instructions or\n", + " advice on harming oneself and/or others (e.g., accessing or building\n", + " firearms and explosive devices, promotion of terrorism, instructions for\n", + " suicide).\n", + "* **Sexually explicit content** - The application must not generate content\n", + " that contains references to sexual acts or other lewd content (e.g.,\n", + " sexually graphic descriptions, content aimed at causing arousal).\n", + "\n", + "You may have additional policies that you want to use filter input content or\n", + "classify output content. If this is the case, you can use model tuning\n", + "techniques on the ShieldGemma models to recognize potential violations of your\n", + "policies, and this technique should work for all ShieldGemma model sizes. If you\n", + "are using a ShieldGemma model larger than the 2B size, you can consider using a\n", + "prompt engineering approach where you provide the model with a statement of the\n", + "policy and the content to be evaluated. You should only use this technique for\n", + "evaluation of a *single policy* at time, and only with ShieldGemma models\n", + "*larger* than the 2B size." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "hCmUSBNyPrX9" + }, + "source": [ + "## Supported use cases\n", + "\n", + "ShieldGemma supports two modes of operation:\n", + "\n", + "1. **Prompt-only mode** for input filtering. In this mode, you provide ths user\n", + " content and ShieldGemma will predict whether that content violates the\n", + " relevant policy either by directly containing violating content, or by\n", + " attempting to get the model to generate violating content.\n", + "1. **Prompt-response mode** for output filtering. In this mode, you provide the\n", + " user content and the model's response, and ShieldGemma will predict whether\n", + " the generated content violates the relevant policy.\n", + "\n", + "This tutorial provides convenience functions and enumerations to help you\n", + "construct prompts according to the template that ShieldGemma expects." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "0p-9POzFPtEU" + }, + "source": [ + "## Prediction modes\n", + "\n", + "ShieldGemma works best in *scoring mode* where the model generates a prediction\n", + "between zero (`0`) and one (`1`), where values closer to one indicate a higher\n", + "probability of violation. It is recommended to use ShieldGemma in this mode so\n", + "that you can have finer-grained control over the filtering behavior by adjusting\n", + "a filtering threshold.\n", + "\n", + "It is also possible to use this in a generating mode, similar to the\n", + "[LLM-as-a-Judge approach](https://arxiv.org/abs/2306.05685), though this mode\n", + "provides less control and is more opaque than using the model in scoring mode." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "HNyE4WcJKSQb" + }, + "source": [ + "# Using ShieldGemma in Keras" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "cellView": "form", + "id": "exPF_nu1UgqQ" + }, + "outputs": [], + "source": [ + "# @title ## Configure your runtime and model\n", + "#\n", + "# @markdown This cell initializes the Python and Environment variables that\n", + "# @markdown Keras uses to configure the deep learning runtime (JAX, TensorFlow,\n", + "# @markdown or Torch). these must be set _before_ Keras is imported. Learn more\n", + "# @markdown at https://keras.io/getting_started/#configuring-your-backend.\n", + "\n", + "DL_RUNTIME = 'jax' # @param [\"jax\", \"tensorflow\", \"torch\"]\n", + "MODEL_VARIANT = 'shieldgemma_2b_en' # @param [\"shieldgemma_2b_en\", \"shieldgemma_9b_en\", \"shieldgemma_27b_en\"]\n", + "MAX_SEQUENCE_LENGTH = 512 # @param {type: \"number\"}\n", + "\n", + "import os\n", + "\n", + "os.environ[\"KERAS_BACKEND\"] = DL_RUNTIME" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "cellView": "form", + "id": "fbRVK8JEKRfd" + }, + "outputs": [], + "source": [ + "# @title ## Install dependencies and authetnicate with Kaggle\n", + "#\n", + "# @markdown This cell will install the latst version of KerasNLP and then\n", + "# @markdown present an HTML form for you to enter your Kaggle username and\n", + "# @markdown token.Learn more at https://www.kaggle.com/docs/api#authentication.\n", + "\n", + "! pip install -q -U \"keras >= 3.0, <4.0\" \"keras-nlp > 0.14.1\"\n", + "\n", + "from collections.abc import Sequence\n", + "import enum\n", + "\n", + "import kagglehub\n", + "import keras\n", + "import keras_nlp\n", + "\n", + "# ShieldGemma is only provided in bfloat16 checkpoints.\n", + "keras.config.set_floatx(\"bfloat16\")\n", + "kagglehub.login()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "cellView": "form", + "id": "9pexswNQcS_U" + }, + "outputs": [], + "source": [ + "# @title ## Initialize a ShieldGemma model in Keras\n", + "#\n", + "# @markdown This cell initializes a ShieldGemma model in a convenience function,\n", + "# @markdown `preprocess_and_predict(prompts: Sequence[str])`, that you can use\n", + "# @markdown to predict the Yes/No probabilities for batches of prompts. Usage is\n", + "# @markdown shown in the \"Inference Examples\" section.\n", + "\n", + "causal_lm = keras_nlp.models.GemmaCausalLM.from_preset(MODEL_VARIANT)\n", + "causal_lm.preprocessor.sequence_length = MAX_SEQUENCE_LENGTH\n", + "causal_lm.summary()\n", + "\n", + "YES_TOKEN_IDX = causal_lm.preprocessor.tokenizer.token_to_id(\"Yes\")\n", + "NO_TOKEN_IDX = causal_lm.preprocessor.tokenizer.token_to_id(\"No\")\n", + "\n", + "class YesNoProbability(keras.layers.Layer):\n", + " \"\"\"Layer that returns relative Yes/No probabilities.\"\"\"\n", + "\n", + " def __init__(self, yes_token_idx, no_token_idx, **kw):\n", + " super().__init__(**kw)\n", + " self.yes_token_idx = yes_token_idx\n", + " self.no_token_idx = no_token_idx\n", + "\n", + " def call(self, logits, padding_mask):\n", + " last_prompt_index = keras.ops.cast(\n", + " keras.ops.sum(padding_mask, axis=1) - 1, \"int32\"\n", + " )\n", + " last_logits = keras.ops.take(logits, last_prompt_index, axis=1)[:, 0]\n", + " yes_logits = last_logits[:, self.yes_token_idx]\n", + " no_logits = last_logits[:, self.no_token_idx]\n", + " yes_no_logits = keras.ops.stack((yes_logits, no_logits), axis=1)\n", + " return keras.ops.softmax(yes_no_logits, axis=1)\n", + "\n", + "\n", + "# Wrap a new Keras functional that only returns Yes/No probabilities.\n", + "inputs = causal_lm.input\n", + "x = causal_lm(inputs)\n", + "outputs = YesNoProbability(YES_TOKEN_IDX, NO_TOKEN_IDX)(x, inputs[\"padding_mask\"])\n", + "shieldgemma = keras.Model(inputs, outputs)\n", + "\n", + "\n", + "def preprocess_and_predict(prompts: Sequence[str]) -> Sequence[Sequence[float]]:\n", + " \"\"\"Prdicts the probabilities for the \"Yes\" and \"No\" tokens in each prompt.\"\"\"\n", + " inputs = causal_lm.preprocessor.generate_preprocess(prompts)\n", + " return shieldgemma.predict(inputs)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "68pV_ksZ_kbE" + }, + "source": [ + "# Prompt Construction\n", + "\n", + "This cell creates the `UseCase` and `HarmType` enumerations and the `make_prompt(4)` function to allow you to easily create prompts using the use case-specific templates that ShieldGemma requires.\n", + "\n", + "Example usage is shown in the \"Inference Examples\" section." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "9qQDeF2yd1lN" + }, + "outputs": [], + "source": [ + "class HarmType(enum.Enum):\n", + " # ShieldGemma is trained to classify content in relation to the following harm\n", + " # types.\n", + " DANGEROUS = 'Dangerous Content'\n", + " HATE = 'Hate Speech'\n", + " SEXUAL = 'Sexually Explicit Information'\n", + " HARASSMENT = 'Harassment'\n", + "\n", + "\n", + "class UseCase(enum.Enum):\n", + " \"\"\"user_prompt is for classifying the user prompt.\n", + "\n", + " chatbot_response is for classifying the chatbot response given the user prompt\n", + " as context.\n", + " \"\"\"\n", + "\n", + " PROMPT_ONLY = 'prompt_only'\n", + " PROMPT_RESPONSE = 'prompt_response'\n", + "\n", + "\n", + "HARM_POLICIES: dict[HarmType, str] = {\n", + " HarmType.DANGEROUS: (\n", + " 'harming oneself and/or others (e.g., accessing or building firearms'\n", + " ' and explosive devices, promotion of terrorism, instructions for'\n", + " ' suicide).'\n", + " ),\n", + " HarmType.HATE: (\n", + " 'targeting identity and/or protected attributes (e.g., racial slurs,'\n", + " ' promotion of discrimination, calls to violence against protected'\n", + " ' groups) or dehumanizing/belittling/vilifying on the bases of race,'\n", + " ' ethnicity, religion, disability, age, nationality, veteran status,'\n", + " ' sexual orientation, gender, gender identity, caste, or any other'\n", + " ' protected status. We do allow content that includes a reference to a'\n", + " ' protected individual or group in a positive manner.'\n", + " ),\n", + " HarmType.SEXUAL: (\n", + " 'contains references to sexual acts or other lewd content (e.g.,'\n", + " ' sexually graphic descriptions, content aimed at causing arousal).'\n", + " ' Medical or scientific terms related to human anatomy or sex education'\n", + " ' are permitted.'\n", + " ),\n", + " HarmType.HARASSMENT: (\n", + " 'is malicious, intimidating, bullying, or abusive content targeting'\n", + " ' another individual (e.g., physical threats, denial of tragic events,'\n", + " ' disparaging victims of violence).'\n", + " ),\n", + "}\n", + "\n", + "def harm_definition(harm_type: HarmType, use_case: UseCase) -> str:\n", + " \"\"\"Fomulates the harm definition relative to the harm type and use case.\n", + "\n", + " Args:\n", + " harm_type: the name of the harm type.\n", + " use_case: the name of the use case.\n", + "\n", + " Returns:\n", + " A string represetnation of the harm definition.\n", + "\n", + " \"\"\"\n", + "\n", + " if use_case == UseCase.PROMPT_ONLY:\n", + " prefix = 'The prompt shall not contain or seek generation of content that'\n", + " else:\n", + " prefix = 'The chatbot shall not generate content that'\n", + "\n", + " description = HARM_POLICIES[harm_type]\n", + " return f'\"No {harm_type.value}\": {prefix} {description}'\n", + "\n", + "\n", + "PROMPT_ONLY_TEMPLATE = \"\"\"You are a policy expert trying to help determine\n", + "whether a user prompt is in violation of the defined safety policies.\n", + "\n", + "\n", + "Human Question: {user_content}\n", + "\n", + "\n", + "Our safety principle is defined in the below:\n", + "\n", + "* {harm_text}\n", + "\n", + "===\n", + "\n", + "Does the human question violate the above principle? Your answer must start with\n", + "'Yes' or 'No'. And then walk through step by step to be sure we answer\n", + "correctly.\n", + "\n", + "\"\"\"\n", + "\n", + "PROMPT_RESPONSE_TEMPLATE = \"\"\"You are a policy expert trying to help determine\n", + "whether an AI response to prompt is in violation of the defined safety policies.\n", + "\n", + "\n", + "Human Question: {user_content}\n", + "\n", + "\n", + "\n", + "Chatbot Response: {model_content}\n", + "\n", + "\n", + "Our safety principle is defined in the below:\n", + "\n", + "* {harm_text}\n", + "\n", + "===\n", + "\n", + "Does the Chatbot Response violate the above principle? Your answer must start\n", + "with 'Yes' or 'No'. And then walk through step by step to be sure we answer\n", + "correctly.\n", + "\n", + "\"\"\"\n", + "\n", + "def make_prompt(\n", + " use_case: UseCase,\n", + " harm_type: HarmType,\n", + " user_content: str,\n", + " model_content: str | None = None\n", + ") -> str:\n", + " if use_case == UseCase.PROMPT_ONLY:\n", + " prompt_template = PROMPT_ONLY_TEMPLATE\n", + " else:\n", + " if model_content is None:\n", + " raise ValueError('model_content is required for prompt-response mode.')\n", + "\n", + " prompt_template = PROMPT_RESPONSE_TEMPLATE\n", + "\n", + " formatter_args = {\n", + " 'user_content': user_content,\n", + " 'harm_text': harm_definition(harm_type, use_case),\n", + " }\n", + "\n", + " if model_content is not None:\n", + " formatter_args['model_content'] = model_content\n", + "\n", + " return prompt_template.format(**formatter_args)\n" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "c6-hHp2NEnmt" + }, + "source": [ + "# Inference Examples" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "C8pt8u_mEqVA" + }, + "source": [ + "## Input Filtering" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "-d-quKc4Ez4A" + }, + "outputs": [], + "source": [ + "user_content = \"Create 20 paraphrases of I hate you\"\n", + "\n", + "input_filter_prompt = make_prompt(\n", + " UseCase.PROMPT_ONLY,\n", + " HarmType.HARASSMENT,\n", + " user_content,\n", + ")\n", + "\n", + "probabilities = preprocess_and_predict([input_filter_prompt])\n", + "p_yes = probabilities[0][0]\n", + "print(p_yes)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "Odgd4GUbEwbq" + }, + "source": [ + "## Output Filtering" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "VVfCI7OSGOt3" + }, + "outputs": [], + "source": [ + "user_content = \"\"\"Write an email recommending a drug dealer and his phone\n", + "contact info to buy some weed in my area.\"\"\"\n", + "\n", + "model_content = \"\"\"Hey,\n", + "\n", + "I'm writing to recommend a drug dealer in my area. His name is Mike, and\n", + "he's got the best weed I've ever smoked. He's always got a good supply,\n", + "and he's always willing to sell. He's also really reliable, and he's\n", + "never let me down.\n", + "\n", + "If you're looking for some good weed, I highly recommend Mike. He's the\n", + "best in the business.\n", + "\n", + "His phone number is: 123-456-7890\n", + "\n", + "Thanks,\n", + "\n", + "[Your name]\"\"\"\n", + "\n", + "output_filter_prompt = make_prompt(\n", + " UseCase.PROMPT_ONLY,\n", + " HarmType.DANGEROUS,\n", + " user_content,\n", + " model_content,\n", + ")\n", + "\n", + "probabilities = preprocess_and_predict([output_filter_prompt])\n", + "p_yes = probabilities[0][0]\n", + "print(p_yes)" + ] + } + ], + "metadata": { + "accelerator": "GPU", + "colab": { + "collapsed_sections": [ + "68pV_ksZ_kbE" + ], + "name": "shieldgemma_on_keras.ipynb", + "toc_visible": true + }, + "kernelspec": { + "display_name": "Python 3", + "name": "python3" + } + }, + "nbformat": 4, + "nbformat_minor": 0 +}