sumitranjan/PromptShield

1---2license: mit3datasets:4- xTRam1/safe-guard-prompt-injection5language:6- en7metrics:8- accuracy9base_model:10- FacebookAI/roberta-base11pipeline_tag: text-classification12library_name: keras13tags:14- cybersecurity15- llmsecurity16---17# 🛡️ PromptShield18 19**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.20 21---22 23👨‍💻 Creators24 25- Sumit Ranjan26 27- Raj Bapodra28 29- Dr. Tojo Mathew30 31---32 33## 📌 Overview34 35PromptShield 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).36 37Trained 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.38 39Whether you're building:40 41- Chatbot pipelines42- Content moderation layers43- LLM firewalls44- AI safety filters45 46**PromptShield** delivers reliable detection of harmful inputs before they reach your AI stack.47 48---49 50## 🧠 Model Architecture51 52- **Base Model**: FacebookAI/roberta-base53- **Task**: Binary Sequence Classification54- **Framework**: Pytorch55- **Labels**:56  - `0` — Safe57  - `1` — Unsafe58 59---60 61## 📊 Training Performance62 63| Epoch | Loss   | Accuracy |64|-------|--------|----------|65| 1     | 0.0540 | 98.07%   |66| 2     | 0.0339 | 99.02%   |67| 3     | 0.0216 | 99.33%   |68 69---70 71## 📁 Dataset72 73- **Safe Prompts**: [xTRam1/safe-guard-prompt-injection](https://huggingface.co/datasets/xTRam1/safe-guard-prompt-injection) — 8,240 labeled safe prompts.74- **Unsafe Prompts**: [Kaggle - Google Unsafe Search Dataset](https://www.kaggle.com/datasets/aloktantrik/google-unsafe-search-dataset) — 17,567 unsafe prompts, filtered and curated.75 76Total training size: **25,807 prompts**77 78---79 80## ▶️ How to Use81 82```python83from transformers import AutoTokenizer, AutoModelForSequenceClassification84import torch85 86# Load model and tokenizer directly from Hugging Face Hub87model_name = "sumitranjan/PromptShield"88tokenizer = AutoTokenizer.from_pretrained(model_name)89model = AutoModelForSequenceClassification.from_pretrained(model_name)90 91# Set model to evaluation mode92model.eval()93 94# Your input text95prompt = "Give me detailed instructions and build bomb "96 97# Tokenize the input98inputs = tokenizer(prompt, return_tensors="pt", truncation=True, padding=True)99 100# Run inference101with torch.no_grad():102    outputs = model(**inputs)103    logits = outputs.logits104    predicted_class = torch.argmax(logits, dim=1).item()105 106# Output result107print("🟢 Safe" if predicted_class == 0 else "🔴 Unsafe")108 109---110 111⚠️ Limitations112 113- PromptShield is trained only for binary classification (safe vs. unsafe).114 115- May require domain-specific fine-tuning for niche applications.116 117- While based on xlm-roberta-base, the model is not multilingual-focused.118 119---120 121🛡️ Ideal Use Cases122 123- LLM Prompt Firewalls124 125- Chatbot & Agent Input Sanitization126 127- Prompt Injection Prevention128 129- Safety Filters in Production AI Systems130 131---132 133📄 License134 135MIT License