---
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.