library_name: transformers
license: mit
language:
- en
tags:
- safety
- guardrail
- text-classification
- distilbert
- llm-guardrails
base_model: distilbert-base-uncased
pipeline_tag: text-classification
Search Query Safety Guard
A DistilBERT-base-uncased guardrail that classifies an AI agent's search query as SAFE or DANGEROUS before it reaches a web search/scrape tool. Designed to sit as a proxy between an agent's search tool and the real search service.
- 0 → SAFE — allow the search
- 1 → DANGEROUS — block and return an error to the agent
Intended use
Wrap a search/scrape tool so every query is classified first. Safe queries pass through to the real service; dangerous queries are blocked and surfaced to the agent as an error. Protected categories include terrorism, CSAM / harm to children, drugs, weapons & explosives, violence, self-harm, hate & extremism, kidnapping & trafficking, fraud & counterfeit, arson, and credential/secret theft (stolen OAuth tokens, leaked API keys, session cookies, etc.).
Cyber policy (intent, not keywords)
This model also backs a cyber-defense agent, so the security boundary is by intent:
| SAFE — let the defender work | DANGEROUS — block |
|---|---|
| "how ransomware works and incident response steps" | "where to buy stolen oauth tokens and session cookies" |
| "how credential stuffing works and how to stop it" | "how to scrape github at scale for accidentally committed private keys" |
| "how reverse shells work and how to detect them" | "where to find leaked api keys and access tokens for sale" |
Rule of thumb: researching how an attack works, or how to detect/prevent/train for it, is SAFE. Stealing, buying, selling, or exfiltrating auth tokens, secrets, credentials, or PII — or targeting a specific person's accounts — is DANGEROUS.
Usage
from transformers import AutoModelForSequenceClassification, AutoTokenizer
import torch
model = AutoModelForSequenceClassification.from_pretrained("urbanspr1nter/search-query-safety-guard")
tokenizer = AutoTokenizer.from_pretrained("urbanspr1nter/search-query-safety-guard")
model.eval()
def predict(query: str) -> dict:
enc = tokenizer(query, return_tensors="pt", max_length=256, truncation=True, padding=True)
with torch.no_grad():
probs = torch.softmax(model(**enc).logits, dim=-1)
pred = int(torch.argmax(probs, dim=-1).item())
return {
"query": query,
"label": model.config.id2label[pred], # "SAFE" or "DANGEROUS"
"safe": pred == 0,
"confidence": float(probs[0, pred].item()),
}
predict("how to cook pasta") # SAFE
predict("where to buy stolen oauth tokens") # DANGEROUS
Evaluation
Trained and evaluated with two held-out sets, both verified disjoint from the training data:
| Set | Accuracy | DANGER F1 | Dangerous leaks (FN) | Safe blocked (FP) |
|---|---|---|---|---|
| eval.jsonl (137 examples) | 99.27% | 0.9937 | 0 | 1 |
| eval2.jsonl (43 examples, fresh) | 100% | 1.0000 | 0 | 0 |
Recall on dangerous queries is 1.000 on both held-out sets — no dangerous query reached the web. A second held-out set was used specifically to guard against memorizing the first.
Training
- Base:
distilbert-base-uncased, 2 labels, max 256 tokens - Optimizer: AdamW, lr 2e-5, batch 8, early stopping on F1
- Data: 2,313 examples (~61% SAFE / ~39% DANGEROUS), grown from 1,973 with a leakage guard rejecting any addition ≥82% similar to an eval query
Limitations
- English only
- Evaluates a single query — no conversation history
- A small dataset cannot cover every edge case; deploy with a confidence threshold and log the grey zone