pip install transformers torch

from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
from torch.nn.functional import softmax

model_id = "meta-llama/Prompt-Guard-86M"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForSequenceClassification.from_pretrained(model_id)
model.eval()

def get_jailbreak_score(text):
    """Check user input for jailbreak attempts."""
    inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=512)
    with torch.no_grad():
        logits = model(**inputs).logits
    probs = softmax(logits, dim=-1)
    return probs[0, 2].item()  # Jailbreak probability

**Classification labels**:
- **BENIGN** (label 0): Normal content
- **INJECTION** (label 1): Embedded instructions in data
- **JAILBREAK** (label 2): Direct override attempts

def filter_user_input(user_message, threshold=0.5):
    """
    Filter user input for jailbreak attempts.

    Returns: (is_safe, score, message)
    """
    score = get_jailbreak_score(user_message)

    if score >= threshold:
        return False, score, "Input blocked: jailbreak attempt"
    else:
        return True, score, "Input safe"

**Common jailbreak patterns detected**:
```python

undefined

def get_indirect_injection_score(text):
    """Check third-party data for embedded instructions."""
    inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=512)
    with torch.no_grad():
        logits = model(**inputs).logits
    probs = softmax(logits, dim=-1)
    # Sum INJECTION + JAILBREAK probabilities
    return (probs[0, 1] + probs[0, 2]).item()

def filter_third_party_data(data, threshold=0.3):
    """
    Filter third-party data (API responses, web scraping, RAG docs).

    Use lower threshold (0.3) for third-party data.
    """
    score = get_indirect_injection_score(data)

    if score >= threshold:
        return False, score, "Data blocked: suspected injection"
    else:
        return True, score, "Data safe"

**Common injection patterns detected**:
```python

undefined

def batch_filter_documents(documents, threshold=0.3, batch_size=32):
    """
    Batch filter documents for prompt injections.

    Args:
        documents: List of document strings
        threshold: Detection threshold (default 0.3)
        batch_size: Batch size for processing

    Returns:
        List of (doc, score, is_safe) tuples
    """
    results = []

    for i in range(0, len(documents), batch_size):
        batch = documents[i:i + batch_size]

        # Tokenize batch
        inputs = tokenizer(
            batch,
            return_tensors="pt",
            padding=True,
            truncation=True,
            max_length=512
        )

        with torch.no_grad():
            logits = model(**inputs).logits

        probs = softmax(logits, dim=-1)
        # Injection scores (labels 1 + 2)
        scores = (probs[:, 1] + probs[:, 2]).tolist()

        for doc, score in zip(batch, scores):
            is_safe = score < threshold
            results.append((doc, score, is_safe))

    return results

undefined

undefined

undefined

undefined

**Solution**: Context-aware filtering with user reputation:
```python
def filter_with_context(text, user_is_trusted):
    score = get_jailbreak_score(text)
    # Higher threshold for trusted users
    threshold = 0.7 if user_is_trusted else 0.5
    return score < threshold

undefined

**Solution**: Sliding window with overlapping chunks:
```python
def score_long_text(text, chunk_size=512, overlap=256):
    """Score long texts with sliding window."""
    tokens = tokenizer.encode(text)
    max_score = 0.0

    for i in range(0, len(tokens), chunk_size - overlap):
        chunk = tokens[i:i + chunk_size]
        chunk_text = tokenizer.decode(chunk)
        score = get_jailbreak_score(chunk_text)
        max_score = max(max_score, score)

    return max_score