| --- |
| license: apache-2.0 |
| base_model: google/gemma-4-E2B-it |
| base_model_relation: finetune |
| library_name: transformers |
| language: |
| - en |
| datasets: |
| - nvidia/Aegis-AI-Content-Safety-Dataset-2.0 |
| - PKU-Alignment/BeaverTails |
| - lmsys/toxic-chat |
| - allenai/wildguardmix |
| - ys-zong/VLGuard |
| - AIML-TUDA/LlavaGuard |
| - yiting/UnsafeBench |
| tags: |
| - gemma-4 |
| - safety |
| - content-moderation |
| - guardrail |
| - shieldgemma |
| - lora |
| - image-text-to-text |
| --- |
| # gemma-shield-e2b |
|
|
| A ShieldGemma-style safety classifier for **text and images**, fine-tuned from `google/gemma-4-E2B-it` |
| with LoRA (merged into the base weights here). One model answers, for a user prompt, a prompt/response |
| pair, or an image, whether it violates a given safety policy. Output is the probability of the single |
| token `Yes` versus `No` at the first generated position. |
|
|
| > Independent fine-tune of Google's Gemma 4 E2B-it. Not a Google release and not affiliated with the |
| > ShieldGemma team; "Gemma" is used only to describe the base model. |
|
|
| ## Policies |
|
|
| Text (ShieldGemma v1 wording): dangerous content, harassment, hate speech, sexually explicit, violence, |
| obscenity, plus a composite "any harm" policy. Image (ShieldGemma 2 wording): sexually explicit, |
| violence/gore, dangerous content, plus a composite. The classifier is policy-conditioned: it also |
| follows policy text it was not trained on (tested with OpenAI Moderation's own eight definitions). |
| Exact prompt templates and policy strings are in `shield_policies.json`. |
|
|
| ## Usage |
|
|
| ```python |
| from transformers import AutoProcessor, Gemma4ForConditionalGeneration |
| import torch |
| m = Gemma4ForConditionalGeneration.from_pretrained("tbtlr/gemma-shield-e2b", dtype=torch.bfloat16).cuda().eval() |
| p = AutoProcessor.from_pretrained("tbtlr/gemma-shield-e2b") |
| yes, no = p.tokenizer.encode("Yes", add_special_tokens=False)[0], p.tokenizer.encode("No", add_special_tokens=False)[0] |
| # Build the prompt from shield_policies.json (template + policy text), run one forward pass and read |
| # the first generated position: logits = m(**inputs, logits_to_keep=1).logits[:, -1] |
| # p_violation = softmax([logits[yes], logits[no]])[0] |
| ``` |
|
|
| Reference scorers, data mappings and the full evaluation suite live in the project repository |
| (`shield/prompt.py`, `shield/score*.py`). |
|
|
| ## Training |
|
|
| - Data: Aegis 2.0, BeaverTails, ToxicChat (benign), WildGuardMix for text; VLGuard, LlavaGuard (SMID), |
| UnsafeBench for images; COCO val2017 as benign image negatives. Source categories are mapped onto the |
| ShieldGemma policies as positive / ambiguous (dropped) / negative. |
| - Recipe: LoRA r=16 on the language model's attention and MLP projections (vision/audio towers frozen), |
| cross-entropy on the answer token, bf16, lr 1e-4 (text, 79k rows) then 5e-5 (image continuation, |
| 23.5k image rows + 20k text replay), one epoch each, on a DGX Spark. |
|
|
| ## Evaluation (Optimal F1 / AU-PRC, composite policy unless noted) |
|
|
| Text: ToxicChat 0.767 / 0.834 (ShieldGemma 2B paper: 0.704 / 0.778), Aegis 2.0 test 0.826 / 0.897, |
| BeaverTails 0.839 / 0.918, WildGuardTest 0.780 / 0.866, XSTest 0.886 / 0.960, OpenAI Moderation |
| 0.787 / 0.856. Image, UnsafeBench test: sexual 0.841 / 0.908, dangerous 0.684 / 0.723, violence |
| 0.652 / 0.687, composite 0.692 / 0.776; LlavaGuard test composite 0.679 / 0.675; VLGuard test |
| composite 0.743 / 0.804. Full tables and error analyses: project README. |
|
|
| A sibling model, `tbtlr/gemma-shield-e2b-trimodal`, adds spoken-request classification at a small cost |
| on OpenAI Moderation (0.761 vs 0.787 F1). |
|
|
| ## License and training data |
|
|
| Weights and code: Apache 2.0, the same license as the base model (`LICENSE` in this repo). The model was |
| trained on the datasets below; two of the text sets are CC BY-NC 4.0. Whether that restriction carries |
| over to model weights is legally unsettled, so commercial users should review the list and make their |
| own assessment. |
|
|
| | dataset | role | declared license | |
| |---|---|---| |
| | Aegis 2.0 (NVIDIA) | text train / eval | CC BY 4.0 | |
| | BeaverTails | text train / eval | CC BY-NC 4.0 | |
| | ToxicChat | text train (benign) / eval | CC BY-NC 4.0 | |
| | WildGuardMix | text train / eval | ODC-BY | |
| | XSTest | text eval | CC BY 4.0 | |
| | OpenAI Moderation eval | text eval | MIT | |
| | VLGuard | image train / eval | MIT | |
| | LlavaGuard (images from SMID) | image train / eval | Apache 2.0 (SMID: CC BY 4.0 original release) | |
| | UnsafeBench | image train / eval | gated data-use agreement (research, education, responsible commercial use) | |
| | COCO val2017 | image train (benign) | CC BY 4.0 annotations; images under their Flickr licenses | |
|
|
| ## Limitations |
|
|
| Binary "any unsafe" scores depend on how far a benchmark's taxonomy is from these policies (hate, |
| political, privacy, fraud and misinformation images/text are out of scope). The dangerous-content image |
| policy was trained with a broader reading (weapons, substance abuse, criminal planning, self-harm imagery) |
| than its text counterpart. English only. Not a substitute for human review in high-stakes settings. |
|
|
