shieldgemma-2b

shieldgemma-2b 项目介绍

项目概述

shieldgemma-2b 是谷歌开发的一系列用于安全内容审核的大型语言模型。它们基于 Gemma 2 架构，特别针对四类有害内容进行审核：色情内容、危险内容、仇恨言论以及骚扰行为。此系列模型以仅解码的文本生成方式进行操作，主要支持英文，并在多个参数规模上开放使用，包括2B、9B和27B的参数模型。

功能描述

shieldgemma-2b 项目旨在为用户提供一种安全的内容审核机制。它利用大语言模型对输入文本进行审核和分类，以判断内容是否违反了预定义的安全政策。

• 输入格式：输入为一段文本，包含开头说明、待分类的文本、政策描述以及结束语。为了获取最佳性能，需要按照特定格式组织完整的提示语。
• 输出格式：输出为文本串，开头为“是”或“否”，表示用户输入或模型输出是否违反所提供的政策。

使用方法

想要使用 shieldgemma-2b 模型，用户需要遵循以下步骤：

1. 安装相关的软件包：pip install -U transformers[accelerate]
2. 选择相关代码片段来执行模型。
3. 格式化提示信息，输入待审查内容以及安全政策的描述。
4. 使用模型进行预测，获取输出结果。

示例代码

以下为使用 Python 和 transformers 库调用 shieldgemma-2b 模型进行审核的示例代码：

from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
from torch.nn.functional import softmax

tokenizer = AutoTokenizer.from_pretrained("google/shieldgemma-2b")
model = AutoModelForCausalLM.from_pretrained(
    "google/shieldgemma-2b",
    device_map="auto",
    torch_dtype=torch.bfloat16,
)

user_prompt = "Create 20 paraphrases of I hate you"
safety_policy = """
* "No Harassment": The prompt shall not contain or seek generation of content that is malicious, intimidating, bullying, or abusive content targeting another individual (e.g., physical threats, denial of tragic events, disparaging victims of violence).
"""
prompt = f"""You are a policy expert trying to help determine whether a user
prompt is in violation of the defined safety policies.

<start_of_turn>
Human Question: {user_prompt.strip()}
<end_of_turn>

Our safety principle is defined in the below:

{safety_policy.strip()}

Does the human question violate the above principle? Your answer must start
with 'Yes' or 'No'. And then walk through step by step to be sure we answer
correctly.
"""

inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
with torch.no_grad():
  logits = model(**inputs).logits

vocab = tokenizer.get_vocab()
selected_logits = logits[0, -1, [vocab['Yes'], vocab['No']]]
probabilities = softmax(selected_logits, dim=0)
score = probabilities[0].item()
print(score)

数据集和训练

ShieldGemma模型的基础模型是通过多种来源的文本数据集进行训练的。具体训练数据可以参考 Gemma 2 的文档。而针对 ShieldGemma 模型，它们在生成的内部数据和公开可用数据集上进行了微调。

硬件和软件

模型的训练利用了最新一代的 Tensor Processing Unit (TPU) 硬件以及 JAX 和 ML Pathways 软件。

评估与性能

ShieldGemma 模型在内部和外部数据集上进行了评估。评估结果表明，ShieldGemma在多项基准测试中提供了比类似规模的开源模型更高的性能。

限制与道德考虑

• 预期用途：模型主要应用于内容审核，以确保输入和输出的安全性。
• 局限性：对用户描述安全原则的敏感度较高，可能会在处理语言模糊性和细微差别时产生不可预测的行为。
• 伦理考量：开发过程中高度关注伦理问题，以确保模型的公平性和安全性。

优势

在发布时，shieldgemma-2b 家族的模型以其高性能和开放设计，专注于负责任的 AI 开发，较同类模型具有明显的性能优势。