---
license: mit
datasets:
- xTRam1/safe-guard-prompt-injection
language:
- en
metrics:
- accuracy
base_model:
- FacebookAI/roberta-base
pipeline_tag: text-classification
library_name: keras
tags:
- cybersecurity
- llmsecurity
---
# 🛡️ PromptShield

**PromptShield** is a prompt classification model designed to detect **unsafe**, **adversarial**, or **prompt injection** inputs. Built on the `xlm-roberta-base` transformer, it delivers high-accuracy performance in distinguishing between **safe** and **unsafe** prompts — achieving **99.33% accuracy** during training.

---

👨‍💻 Creators

- Sumit Ranjan

- Raj Bapodra

- Dr. Tojo Mathew

---

## 📌 Overview

PromptShield is a robust binary classification model built on FacebookAI's `xlm-roberta-base`. Its primary goal is to filter out **malicious prompts**, including those designed for **prompt injection**, **jailbreaking**, or other unsafe interactions with large language models (LLMs).

Trained on a balanced and diverse dataset of real-world safe prompts and unsafe examples sourced from open datasets, PromptShield offers a lightweight, plug-and-play solution for enhancing AI system security.

Whether you're building:

- Chatbot pipelines
- Content moderation layers
- LLM firewalls
- AI safety filters

**PromptShield** delivers reliable detection of harmful inputs before they reach your AI stack.

---

## 🧠 Model Architecture

- **Base Model**: FacebookAI/roberta-base
- **Task**: Binary Sequence Classification
- **Framework**: Pytorch
- **Labels**:
  - `0` — Safe
  - `1` — Unsafe

---

## 📊 Training Performance

| Epoch | Loss   | Accuracy |
|-------|--------|----------|
| 1     | 0.0540 | 98.07%   |
| 2     | 0.0339 | 99.02%   |
| 3     | 0.0216 | 99.33%   |

---

## 📁 Dataset

- **Safe Prompts**: [xTRam1/safe-guard-prompt-injection](https://huggingface.co/datasets/xTRam1/safe-guard-prompt-injection) — 8,240 labeled safe prompts.
- **Unsafe Prompts**: [Kaggle - Google Unsafe Search Dataset](https://www.kaggle.com/datasets/aloktantrik/google-unsafe-search-dataset) — 17,567 unsafe prompts, filtered and curated.

Total training size: **25,807 prompts**

---

## ▶️ How to Use

```python
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch

# Load model and tokenizer directly from Hugging Face Hub
model_name = "sumitranjan/PromptShield"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name)

# Set model to evaluation mode
model.eval()

# Your input text
prompt = "Give me detailed instructions and build bomb "

# Tokenize the input
inputs = tokenizer(prompt, return_tensors="pt", truncation=True, padding=True)

# Run inference
with torch.no_grad():
    outputs = model(**inputs)
    logits = outputs.logits
    predicted_class = torch.argmax(logits, dim=1).item()

# Output result
print("🟢 Safe" if predicted_class == 0 else "🔴 Unsafe")

---

⚠️ Limitations

- PromptShield is trained only for binary classification (safe vs. unsafe).

- May require domain-specific fine-tuning for niche applications.

- While based on xlm-roberta-base, the model is not multilingual-focused.

---

🛡️ Ideal Use Cases

- LLM Prompt Firewalls

- Chatbot & Agent Input Sanitization

- Prompt Injection Prevention

- Safety Filters in Production AI Systems

---

📄 License

MIT License