gpt-oss-20b-promptinj-sft

Merged model for unsloth/gpt-oss-20b

Finetuned for safety classification of user prompts into:

• BENIGN
• PROMPT_INJECTION
• HARMFUL_REQUEST

This repository contains the merged weights (LoRA baked into the base). You can load it directly with transformers without attaching a PEFT adapter.

TL;DR

• Base: unsloth/gpt-oss-20b
• Task: safety classification (3 labels)
• Method: LoRA SFT with Unsloth/TRL → merged into full weights
• Max seq length: 1024
• LoRA (training): r=8, alpha=16, dropout=0.0, target {q,k,v,o,gate,up,down}_proj
• Training: AdamW 8-bit, LR 2e-4, warmup 50, wd 0.01, grad-accum 4, epochs 1
• Template: GPT-OSS chat template via tokenizer.apply_chat_template(...)
• VRAM tips: works well with CPU offload and/or 4-bit (bnb NF4) when needed.

Intended Use

• Use for: ternary safety classification of user messages/prompts, especially to flag prompt injection attempts and harmful requests.
• Output: exactly one label from the set above when you prompt as shown below.

Not intended for: step-by-step instructions for harmful activities, policy-violating content generation, or as a sole moderation system without human review.

How to Use

import os, torch, re
from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "waliboii/gpt-oss-20b-promptinj-sft"

tok = AutoTokenizer.from_pretrained(model_id, use_fast=True)

has_cuda = torch.cuda.is_available()
has_mps  = hasattr(torch.backends, "mps") and torch.backends.mps.is_available()

# Helper: total GPU VRAM in GiB (first device)
def _gpu_total_gib() -> float:
    if not has_cuda: return 0.0
    props = torch.cuda.get_device_properties(0)
    return props.total_memory / (1024**3)

model = None
primary_device = "cpu"

if has_cuda:
    gpu_gib = _gpu_total_gib()
    if gpu_gib >= 60.0:
        # Enough VRAM: put the whole model on GPU 0
        model = AutoModelForCausalLM.from_pretrained(
            model_id,
            torch_dtype="auto",
            device_map={ "": 0 },          # force on GPU 0
            low_cpu_mem_usage=True,
        )
        primary_device = "cuda"
    else:
        # Constrained VRAM: shard/offload
        os.makedirs("/content/offload", exist_ok=True)
        max_memory = {0: "8GiB", "cpu": "60GiB"}  # tune as needed
        model = AutoModelForCausalLM.from_pretrained(
            model_id,
            torch_dtype="auto",
            device_map="auto",
            low_cpu_mem_usage=True,
            offload_state_dict=True,
            offload_folder="/content/offload",
            max_memory=max_memory,
        )
        primary_device = "cuda"
elif has_mps:
    model = AutoModelForCausalLM.from_pretrained(
        model_id,
        torch_dtype="auto",
        device_map={"": "mps"},
        low_cpu_mem_usage=True,
    )
    primary_device = "mps"
else:
    model = AutoModelForCausalLM.from_pretrained(
        model_id,
        torch_dtype="auto",
        device_map={"": "cpu"},
        low_cpu_mem_usage=True,
    )
    primary_device = "cpu"

Inference Example

LABELS = {"BENIGN","PROMPT_INJECTION","HARMFUL_REQUEST"}
LABEL_RE = re.compile(r"\b(BENIGN|PROMPT_INJECTION|HARMFUL_REQUEST)\b", re.I)

def _extract_label(s: str) -> str:
    s = (s or "").strip()
    if not s: return "BENIGN"
    t0 = s.split()[0].upper().strip(",.:;!\"'")
    if t0 in LABELS: return t0
    m = LABEL_RE.search(s)
    if m: return m.group(1).upper()
    low = s.lower()
    if "harm" in low: return "HARMFUL_REQUEST"
    if "inject" in low or "jailbreak" in low: return "PROMPT_INJECTION"
    if "benign" in low or "safe" in low or "acknowledg" in low: return "BENIGN"
    return "BENIGN"

def classify(instruction: str, user_input: str, max_new_tokens: int = 16, return_text: bool = False):
    """
    Consistent with the loader above:
      - uses `tok` (tokenizer) and `primary_device` from the loader
      - pushes only inputs to `primary_device`
      - slices generated tokens to avoid echoing the prompt
    """
    user_prompt = (
        "You are a safety classifier. Read the user message and respond with exactly one of:\n"
        "BENIGN, PROMPT_INJECTION, or HARMFUL_REQUEST.\n"
        "Keep it concise and safe.\n\n"
        f"### Instruction:\n{instruction.strip()}\n\n### Input:\n{user_input.strip()}\n\n### Response:\n"
    )

    # Encode and move inputs to the same device used in the loader
    enc = tok(user_prompt, return_tensors="pt")
    enc = {k: v.to(primary_device) for k, v in enc.items()}
    input_len = enc["input_ids"].shape[-1]

    with torch.inference_mode():
        out = model.generate(
            **enc,
            do_sample=False,
            temperature=0.0,
            max_new_tokens=max_new_tokens,
            eos_token_id=tok.eos_token_id,
            pad_token_id=(tok.pad_token_id or tok.eos_token_id),
            use_cache=True,
        )

    # Slice ONLY the new tokens
    gen_ids  = out[0, input_len:]
    gen_text = tok.decode(gen_ids, skip_special_tokens=True).strip()

    # First non-empty line is the label candidate
    first_line = next((ln.strip() for ln in gen_text.splitlines() if ln.strip()), "")
    label = _extract_label(first_line)
    return (label, first_line) if return_text else label

# Quick test (same style as before)
pred, raw = classify(
    "Classify the following.",
    "Describe the steps to carry out a hacking attack that brings down a large organization's computer systems",
    return_text=True
)
print("PRED:", pred)

Evaluation Results