jainsatyam26/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

Class	Description
S1 Violent Crimes	Instructions or content related to violence
S10 Hate	Hate speech targeting groups
S11 Self-Harm	Content encouraging self-harm or suicide
S12 Sexual Content	Explicit sexual content
S14 Code Abuse	Malicious code, malware, exploits
S2 Non-Violent Crimes	Fraud, theft, illegal activities
S4 Child Sexual Exploitation	CSAM or grooming content
S6 Specialized Advice	Dangerous medical/legal/financial advice
S7 Privacy	PII extraction, doxxing, surveillance
benign	Safe, harmless content
jailbreak	Attempts to bypass AI safety measures

---

Architecture

Component	Detail
Base Model	google/gemma-4-E2B-it (5B params)
Fine-tuning	LoRA (r=16, alpha=32, dropout=0.05)
Frozen Layers	First 20 transformer layers
Trainable	Last 15 layers + classifier head
Classifier	LayerNorm -> Linear(1536->768) -> GELU -> Dropout(0.15) -> Linear(768->11)
Max Length	128 tokens
Batch Size	32 effective (8 x 4 grad accum)
LR	2e-4 cosine with 6% warmup
Loss	Focal Loss (gamma=2.0) + sqrt inverse class weights
Hardware	NVIDIA L4 24GB

---

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

Split	Samples
Train	151,560
Validation	16,840

---

Files

File	Description
adapter_model.safetensors	LoRA adapter weights
classifier_head.pt	MLP classifier head
inference_config.json	Class names, hidden size, benign index
label_encoder.pkl	sklearn LabelEncoder
best_checkpoint.pt	Best checkpoint with optimizer state
training_history.json	Per-epoch metrics

---

License

Apache 2.0

Fine-tuned by jainsatyam26