shieldgemma-2b - 结合安全审核与文本生成技术，实现合规与伦理保障

shieldgemma-2b 项目介绍

项目概述

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

功能描述

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

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

使用方法

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

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

示例代码

以下为使用 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)