jainsatyam26/gemma4-guardrail-v5

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"))

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