Akash-Sakala's picture
Add model card with metrics and dataset info
a7ef4a8 verified
|
Raw
History Blame
3.84 kB
metadata
language:
  - en
license: apache-2.0
base_model: distilbert/distilbert-base-uncased
tags:
  - text-classification
  - phishing-detection
  - knowledge-distillation
  - distilbert
datasets:
  - Akash-Sakala/phishing-site-classification
metrics:
  - accuracy
  - f1
  - precision
  - recall
model-index:
  - name: bert-phishing-classifier_student
    results:
      - task:
          type: text-classification
          name: Text Classification
        dataset:
          name: Akash-Sakala/phishing-site-classification
          type: Akash-Sakala/phishing-site-classification
          split: test
        metrics:
          - type: accuracy
            value: 0.9601
            name: Accuracy
          - type: f1
            value: 0.9595
            name: F1
          - type: precision
            value: 0.971
            name: Precision
          - type: recall
            value: 0.9483
            name: Recall

DistilBERT Phishing Site Classifier (Student)

A 4-layer DistilBERT trained via knowledge distillation from a fine-tuned BERT teacher (Akash-Sakala/bert-phishing-classifier_teacher) for binary phishing site URL classification.

Model Details

Property Value
Base model distilbert/distilbert-base-uncased
Architecture DistilBertForSequenceClassification
Layers 4 (distilled from 12-layer BERT teacher)
Attention heads 8
Task Binary classification (phishing / benign)
Parameters ~52M

Training — Distillation Setup

Hyperparameter Value
Temperature 3.0
Alpha (KL weight) 0.6
Hard label weight 0.4
Learning rate 2e-5
Batch size 64
Epochs 4
Warmup steps 10% of total steps
Weight decay 0.01
Optimizer AdamW
Scheduler Linear with warmup
Mixed precision fp16 (torch.amp)

Loss Function

Combined KL divergence (soft targets) + Cross-Entropy (hard labels):

loss = alpha * KL(student_soft || teacher_soft) * T^2 + (1 - alpha) * CrossEntropy(student, labels)

Test Set Results

Model Accuracy Precision Recall F1
BERT Teacher 0.8971 0.9136 0.8763 0.8945
DistilBERT Student 0.9601 0.9710 0.9483 0.9595

The student outperforms the teacher across all metrics while being smaller and faster.

Dataset

Usage

from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch

tokenizer = AutoTokenizer.from_pretrained('Akash-Sakala/bert-phishing-classifier_student')
model = AutoModelForSequenceClassification.from_pretrained('Akash-Sakala/bert-phishing-classifier_student')

url = 'http://suspicious-login.verify-account.com/secure'
inputs = tokenizer(url, return_tensors='pt', truncation=True, padding='max_length')

with torch.no_grad():
    logits = model(**inputs).logits

pred = torch.argmax(logits, dim=1).item()
print('Phishing' if pred == 1 else 'Benign')

Teacher Model

Akash-Sakala/bert-phishing-classifier_teacher