LLM prompt detection: four baseline architectures
We are comparing how much model complexity an input-side guardrail needs. This repository contains a TF-IDF MLP, a sentence CNN, a bidirectional LSTM, and a fully fine-tuned DistilBERT. The project combines literature and dataset review with data acquisition, preparation, training, and evaluation.
The implemented task is prompt harmfulness classification. Label 0 means benign and label 1 means harmful, with a fixed probability threshold of 0.5. This experiment does not establish context-dependent prompt-injection detection.
Code and report provenance: https://github.com/Deep-Axe/llm-prompt-detection
Data and training
All four models use the same 100,000 unique WildJailbreak training prompts, 39,502 validation prompts, and 38,911 test prompts. The official evaluation adds 2,000 harmful and 210 benign prompts. Inputs exclude model completions. We use upstream vanilla/adversarial benign and harmful data types as labels. Normalized duplicates and label conflicts are removed; variants of the same underlying vanilla request stay together in the main splits. Official evaluation does not expose those identifiers, so only exact overlap can be checked there.
Each architecture trains for three epochs, seed 4442, selecting the checkpoint with highest validation F1. MLP uses 32,768 hashed word/character TF-IDF features and a 64-unit hidden layer. CNN and BiLSTM use a training-only 20,000-term vocabulary and learned 64-dimensional embeddings. All models have a 256-token budget, but regex and WordPiece tokenization differ. MLP trains on local CPU; CNN and BiLSTM on an RTX 3060. DistilBERT trains on one Kaggle T4 with batch 32, FP16, gradient checkpointing, and AdamW at 2e-5. All 66,955,010 parameters are trainable, with encoder gradients and changed weights verified.
Results
| Model | Main test F1 | Official harmful recall | Official benign acceptance |
|---|---|---|---|
| MLP | 0.9355 | 0.8000 | 0.5857 |
| CNN | 0.9290 | 0.8585 | 0.4476 |
| BiLSTM | 0.9423 | 0.8570 | 0.5571 |
| DistilBERT | 0.9792 | 0.9055 | 0.6000 |
DistilBERT leads the main test, but falsely flags 40% of official benign prompts. These single-seed baselines are not ready-made production guardrails. The data are largely synthetic, related tactics can cross request groups, and truncation can remove relevant text. Bootstrap intervals describe test sampling uncertainty, not training-seed variation. Full metrics, histories, hashes, and model settings are in each folder's training_results.json and metadata.json. Checkpoints were reloaded and predictions checked against saved evaluation outputs.
Load DistilBERT
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
repo = "DeeAxe/llm-prompt-detection"
tokenizer = AutoTokenizer.from_pretrained(repo, subfolder="distilbert")
model = AutoModelForSequenceClassification.from_pretrained(repo, subfolder="distilbert").eval()
inputs = tokenizer("Explain how rainbows form.", return_tensors="pt",
truncation=True, max_length=256)
with torch.inference_mode():
harmful_probability = model(**inputs).logits.softmax(-1)[0, 1].item()
Load any of the four models
Use Python 3.10 and install requirements.txt in a separate environment. The project's inference helper accepts this repository ID directly and fetches only the selected architecture's files:
python src/predict_large_models.py --repo-id DeeAxe/llm-prompt-detection --model cnn --text "Explain how rainbows form."
From the source repository, Python usage is:
from src.predict_large_models import PromptClassifier
classifier = PromptClassifier.from_hub("DeeAxe/llm-prompt-detection", architecture="cnn")
probabilities = classifier.probabilities(["Explain how rainbows form."])
Downloads are cached. Inference executes on the local CPU by default; select
--device cuda to use a local GPU. --revision COMMIT_SHA pins a release and
--offline reuses an existing cached version. Output includes the resolved Hub
revision. This model repository distributes checkpoints; no hosted inference
endpoint is deployed.
You can also download the repository and run its bundled helper:
from huggingface_hub import snapshot_download
folder = snapshot_download("DeeAxe/llm-prompt-detection")
print(folder)
python /path/to/snapshot/predict.py --root /path/to/snapshot --model cnn --text "Explain how rainbows form."
Model choices are mlp, cnn, bilstm, and distilbert; the default device is CPU. MLP's joblib checkpoint should only be loaded from a trusted source. The helper reproduces the saved vocabulary, token limits, and training-fitted TF-IDF. CNN/BiLSTM state dictionaries load with torch.load(..., weights_only=True).
Attribution
WildJailbreak: Jiang et al., WildTeaming at Scale (2024), ODC-BY-1.0, revision 5ddc12a7894f842b0619b8e1c7ee496b198af009. DistilBERT base: distilbert/distilbert-base-uncased, Apache-2.0, revision 12040accade4e8a0f71eabdb258fecc2e7e948be. See LICENSE.md for file-specific licenses. Prompt data and individual predictions are excluded; this repository distributes checkpoints and aggregate results.