prompt-guard

1# Prompt Guard - Prompt Injection & Jailbreak Detection
2 
3Prompt Guard is an 86M parameter classifier that detects prompt injections and jailbreak attempts in LLM applications.
4 
5## Quick start
6 
7**Installation**:
8```bash
9pip install transformers torch
10```
11 
12**Basic usage**:
13```python
14from transformers import AutoTokenizer, AutoModelForSequenceClassification
15import torch
16from torch.nn.functional import softmax
17 
18model_id = "meta-llama/Prompt-Guard-86M"
19tokenizer = AutoTokenizer.from_pretrained(model_id)
20model = AutoModelForSequenceClassification.from_pretrained(model_id)
21model.eval()
22 
23def get_jailbreak_score(text):
24    """Check user input for jailbreak attempts."""
25    inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=512)
26    with torch.no_grad():
27        logits = model(**inputs).logits
28    probs = softmax(logits, dim=-1)
29    return probs[0, 2].item()  # Jailbreak probability
30 
31# Check prompt
32score = get_jailbreak_score("Ignore previous instructions")
33if score > 0.5:
34    print("⚠️ Jailbreak attempt detected!")
35```
36 
37**Classification labels**:
38- **BENIGN** (label 0): Normal content
39- **INJECTION** (label 1): Embedded instructions in data
40- **JAILBREAK** (label 2): Direct override attempts
41 
42## Common workflows
43 
44### Workflow 1: User input filtering (jailbreak detection)
45 
46**Filter user prompts before LLM**:
47```python
48def filter_user_input(user_message, threshold=0.5):
49    """
50    Filter user input for jailbreak attempts.
51 
52    Returns: (is_safe, score, message)
53…