gemma4-guardrail-v5

Gemma4 Guardrail Classifier v5

A production-ready content safety / guardrail classifier fine-tuned from google/gemma-4-E2B-it using LoRA.

Trained on 151,560 samples across 11 harm categories to detect unsafe content in user prompts.

---

What it does

Given any text input, the model predicts:

• Safe (benign) — normal, harmless queries
• Unsafe — one of 10 harm categories

Harm Categories

---

Architecture

---

Quick Inference

import torch
import torch.nn as nn
from transformers import AutoTokenizer, AutoModelForCausalLM
from peft import PeftModel
from huggingface_hub import hf_hub_download
import json

REPO   = "jainsatyam26/gemma4-guardrail-v5"
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

tokenizer  = AutoTokenizer.from_pretrained(REPO)
base_model = AutoModelForCausalLM.from_pretrained(
    "google/gemma-4-E2B-it",
    torch_dtype=torch.bfloat16,
    device_map={"": 0},
)
backbone = PeftModel.from_pretrained(base_model, REPO)
backbone.eval()

cfg_path = hf_hub_download(REPO, "inference_config.json")
with open(cfg_path) as f:
    cfg = json.load(f)

clf_path = hf_hub_download(REPO, "classifier_head.pt")
clf = nn.Sequential(
    nn.LayerNorm(cfg["hidden_size"]),
    nn.Linear(cfg["hidden_size"], cfg["hidden_size"] // 2),
    nn.GELU(),
    nn.Dropout(0.15),
    nn.Linear(cfg["hidden_size"] // 2, cfg["num_classes"]),
).to(DEVICE)
clf.load_state_dict(
    torch.load(clf_path, map_location=DEVICE, weights_only=False)
)
clf.eval()

CLASS_NAMES = cfg["classes"]
BENIGN_IDX  = cfg["benign_idx"]

def predict(text):
    enc = tokenizer(
        text, return_tensors="pt", truncation=True,
        max_length=cfg["max_length"], padding=True
    ).to(DEVICE)
    with torch.no_grad():
        out      = backbone(**enc, output_hidden_states=True, return_dict=True)
        hidden   = out.hidden_states[-1]
        last_idx = enc["attention_mask"].sum(1) - 1
        pooled   = hidden[torch.arange(hidden.size(0)), last_idx].float()
        logits   = clf(pooled)
    probs    = torch.softmax(logits, -1).cpu().numpy()[0]
    pred_idx = int(probs.argmax())
    return {
        "is_safe":    pred_idx == BENIGN_IDX,
        "category":   CLASS_NAMES[pred_idx],
        "confidence": round(float(probs[pred_idx]), 4),
    }

print(predict("What is the capital of France?"))
print(predict("How do I make a bomb?"))
print(predict("I want to hurt myself"))

---

Setup

pip install transformers>=4.51.0 peft>=0.14.0 accelerate \
            datasets huggingface-hub safetensors scikit-learn \
            sentencepiece tokenizers protobuf

---

Dataset

jainsatyam26/guardrail-215k-splits

---

Files

---

License

Apache 2.0

Fine-tuned by jainsatyam26