najlajj453/malicious-url-distilbert

Malicious URL Detection — DistilBERT

A fine-tuned DistilBERT classifier that labels a URL as one of four classes — benign, phishing, malware, or defacement — directly from the raw URL string, with no handcrafted feature engineering.

Full training/evaluation code, unit tests, a Gradio demo, and the technical report live in the companion GitHub repository: https://github.com/najla-jaashan/GP-Malicious-URL-Detection

Quick usage

Because the label mapping is embedded in config.json (id2label / label2id), this model works directly with the standard pipeline API — no custom loading code required:

from transformers import pipeline

classifier = pipeline("text-classification", model="najla-jaashan/malicious-url-distilbert")
classifier("http://secure-login.paypa1-account.example/verify")
# -> [{'label': 'phishing', 'score': 0.97...}]

Or load the model and tokenizer directly:

from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch

tokenizer = AutoTokenizer.from_pretrained("najla-jaashan/malicious-url-distilbert")
model = AutoModelForSequenceClassification.from_pretrained("najla-jaashan/malicious-url-distilbert")

inputs = tokenizer("http://192.168.1.1/wp-admin/setup-config.php", return_tensors="pt")
with torch.no_grad():
    probs = torch.softmax(model(**inputs).logits, dim=1)[0]
pred_id = int(probs.argmax())
print(model.config.id2label[pred_id], float(probs[pred_id]))

Label mapping

id	label
0	benign
1	defacement
2	malware
3	phishing

Model details

• Base model: distilbert-base-uncased (6 layers, 66M parameters)
• Task: 4-class sequence classification
• Input: raw URL string, tokenized to a max of 64 sub-word tokens
• Training data: ~651K labeled URLs (the Kaggle "Malicious URLs Dataset"), split 70% train / 15% validation / 15% test, stratified by class
• Optimizer: AdamW, learning rate 2e-5, weight decay 0.01, 3 epochs

Evaluation

Reference results for this training configuration (see the GitHub repo's docs/TECHNICAL_REPORT.md for the full breakdown and reproduction steps):

Class	Precision	Recall	F1-Score
Benign	0.9922	0.9944	0.9933
Phishing	0.9641	0.9596	0.9619
Malware	0.9908	0.9705	0.9805
Defacement	0.9974	0.9991	0.9982
Macro Avg	0.9861	0.9809	0.9835

Overall test accuracy: 98.89% (held-out, in-distribution test split).

These are the reference numbers for this training configuration. Run python -m src.evaluate in the GitHub repo against this checkpoint to confirm the exact metrics for this specific set of weights, and update this table if they differ.

Intended use & limitations

• This is a research prototype, not a production security product. A prediction is one signal, not a guarantee that a URL is safe — use it within a layered security workflow, not as the sole gate.
• Reported accuracy is in-distribution (same source dataset, held-out split). Cross-dataset generalization and robustness to adversarial evasions (homoglyphs, typosquatting, subdomain padding) are evaluated separately in the GitHub repo (src/cross_eval.py, src/adversarial.py) — check there before relying on this model against out-of-distribution or adversarial traffic.
• The repo also ships a calibrated confidence threshold: predictions below it are better treated as "needs review" rather than a hard label (see src/calibrate.py / src/predict.py).

Authorship & credit

Engineering, fine-tuning run, and productionization: Najla Jaashan — the modular training/evaluation pipeline, calibration, cross-dataset and adversarial evaluation, and this packaged checkpoint.