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.

Underlying modeling approach and dataset choice originate with the research team at the Department of Computer Engineering, Imam Abdulrahman Bin Faisal University, Dammam (see the GitHub repo's docs/paper.pdf for the original paper).