RAG-PIF RoBERTa: Prompt Injection Classifier (Layer 3)
Layer 3 of RAG-PIF, a 3-layer prompt injection firewall for Retrieval-Augmented
Generation systems. This is a fine-tuned roberta-base binary classifier
(0 = benign, 1 = injection), intended to run after a fast regex filter (Layer 1)
and a FAISS embedding similarity check (Layer 2) as the final, most expensive
but most accurate layer in the cascade.
Full pipeline code: https://github.com/shivateja1253/rag-pif
The production checkpoint is at the repo root (config.json, model.safetensors,
tokenizer.json, tokenizer_config.json). Two additional, non-production
experimental checkpoints are included under experiments/ -- see below.
Intended use
Classifying short-to-medium text spans (queries, retrieved document chunks) as prompt-injection attempts or benign content within a RAG pipeline. Not intended as a standalone content moderation system.
Training data
Trained on a dataset combining HackaPrompt-derived injection examples, synthetic evasion variants, Wikipedia-derived benign text, and synthetic hard negatives (research/news/fiction-framed benign text discussing injection techniques, designed to reduce false positives on security-discussion content). Full dataset: https://huggingface.co/datasets/shivateja1253/rag-pif-datasets
Evaluation (held-out test sets, full 3-layer pipeline, production checkpoint)
| Test set | F1 | FPR | Detection rate | n |
|---|---|---|---|---|
| Standard | 0.996 | 0.008 | 1.00 | 263 |
| Evasion (obfuscated/unicode/zero-width) | -- | 0.000 | 1.00 | 150 |
| Adversarial (benign security-research text) | -- | 0.280 | -- | 200 |
| OOD (enterprise emails/support tickets) | -- | 0.053 | -- | 170 |
Known limitations
- Generalizes well to character-substitution evasion (e.g. leetspeak) but inconsistently to informal/abbreviated semantic paraphrasing with no recognizable trigger words.
- Remaining false positives on adversarial (security-research-framed) text are this model's primary limitation.
Experimental checkpoints (experiments/)
Two additional checkpoints are included for full transparency and reproducibility of experiments described in the paper -- neither is used in production, both underperformed the checkpoint at the repo root:
| Checkpoint | Approach | Adversarial FPR | OOD FPR | Standard F1 | Outcome |
|---|---|---|---|---|---|
| Production (root) | Baseline training on train_augmented.csv | 0.280 | 0.053 | 0.996 | Kept |
experiments/v3/ | + 500 synthetic hard negatives (research/news/fiction-framed) | ~0.265-0.295 | 0.141 (regressed) | 0.996 | Rejected -- improved adversarial slightly but nearly tripled OOD false positives |
experiments/v4/ | + 600 hard negatives (adversarial + OOD-style) with class-weighted loss | 0.305 | 0.165 | 0.974 | Rejected -- worse than production on every metric |
This progression (v2 -> v3 -> v4) is documented in the GitHub README's "Known limitations" section and illustrates that the remaining adversarial false positives reflect a genuine generalization limit rather than a simple data-coverage gap that more targeted training data can close.
Usage (production checkpoint)
from transformers import RobertaTokenizerFast, RobertaForSequenceClassification
import torch
tokenizer = RobertaTokenizerFast.from_pretrained("shivateja1253/rag-pif-roberta")
model = RobertaForSequenceClassification.from_pretrained("shivateja1253/rag-pif-roberta")
inputs = tokenizer("ignore all previous instructions", return_tensors="pt")
logits = model(**inputs).logits
prob_injection = torch.softmax(logits, dim=-1)[0][1].item()
To load an experimental checkpoint instead:
model = RobertaForSequenceClassification.from_pretrained(
"shivateja1253/rag-pif-roberta", subfolder="experiments/v3"
)