--- datasets: - ealvaradob/phishing-dataset - cybersectony/PhishingEmailDetectionv2.0 - zefang-liu/phishing-email-dataset - kmack/Phishing_urls language: - en pipeline_tag: feature-extraction tags: - transformers - pytorch - distilbert - text-classification - feature-extraction - security - cybersecurity - phishing - malware - url - tiny - lightweight license: apache-2.0 --- This is a lightweight model utilizing the DistilBERT architecture, designed to produce high-quality embeddings for text containing URLs. Despite utilizing the DistilBERT architecture, urlbert-tiny-v5 was not trained via knowledge distillation and is not a fine-tune of the original DistilBERT. Instead, the model was trained on MLM, text generation, token classification, and multi-class classification tasks. ### Key Specifications - Architecture: DistilBERT (6 layers, 768 hidden dimensions, 12 attention heads) - Parameters: ~58.2M - Context Window: 512 tokens - Vocabulary Size: 19,996 - Tensor type: F32 ### Here is a minimal example showing how to extract embeddings from text containing URLs: ``` import torch from transformers import AutoTokenizer, AutoModel model_name = "CrabInHoney/urlbert-tiny-v5" tokenizer = AutoTokenizer.from_pretrained(model_name) model = AutoModel.from_pretrained(model_name) text = "Check that model: https://huggingface.co/CrabInHoney/urlbert-tiny-v5" inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=512) with torch.no_grad(): outputs = model(**inputs) embedding = outputs.last_hidden_state[:, 0, :] print(f"Embedding Shape: {embedding.shape}") print(f"First 5 values: {embedding[0, :5]}") ``` Output: ``` Embedding Shape: torch.Size([1, 768]) First 5 values: tensor([-0.0206, -0.0150, -0.0403, 0.0814, 0.0638]) ``` ### Given that urlbert-tiny-v5 generates high-quality embeddings suitable for classification "out-of-the-box," we decided not to release separate base and fine-tuned versions. Instead, only classification heads were trained for specific datasets, while the encoder weights remained frozen during the process. There are 7 trained heads available in the `heads/` directory of this repository. ### Benchmark Results The following table shows the performance of these heads on their respective test sets: | Model Head File (`.safetensors`) | Dataset Source | Task Type | Samples | Accuracy | Macro F1 | | :--- | :--- | :--- | :--- | :--- | :--- | | `MSMalicious-URLs-dataset_head` | [Kaggle: MS Malicious URLs](https://www.kaggle.com/datasets/sid321axn/malicious-urls-dataset) | **4-Class:** (Benign, Defacement, Phishing, Malware) | 651,191 | **99.82%** | 0.9965 | | `cyPhishing-Email-Detection_head` | [HF: Cybersectony Phishing v2.0](https://huggingface.co/datasets/cybersectony/PhishingEmailDetectionv2.0) | **4-Class:** (Legit/Phish Email, Legit/Phish URL) | 200,000 | **99.69%** | 0.9914 | | `PSSpam-Email-Classification_head` | [Kaggle: Email Spam Classification](https://www.kaggle.com/datasets/purusinghvi/email-spam-classification-dataset) | **Binary:** (Legit vs Spam Email) | 83,448 | **99.10%** | 0.9909 | | `zlphishing-email-dataset_head` | [HF: ZL Phishing Email](https://huggingface.co/datasets/zefang-liu/phishing-email-dataset) | **Binary:** (Safe vs Phish Email) | 18,634 | **97.98%** | 0.9790 | | `eaphishing-dataset_head` | [HF: EA Phishing (Combined)](https://huggingface.co/datasets/ealvaradob/phishing-dataset) | **Binary:** (Safe vs Phishing) | 77,677 | **96.67%** | 0.9660 | | `kmPhishing-urls_head` | [HF: KMack Phishing URLs](https://huggingface.co/datasets/kmack/Phishing_urls) | **Binary:** (Safe vs Phishing URL) | 708,820 | **89.51%** | 0.8948 | | `annotationGenHead` | *Unpublished Dataset* | *Annotation Generation* | - | - | - | ### Inference Example This script loads the base model and **all** available heads to analyze a URL/text against every dataset simultaneously. ``` import torch, torch.nn as nn, torch.nn.functional as F from transformers import AutoTokenizer, AutoModel, AutoModelForCausalLM, EncoderDecoderModel, BertConfig from huggingface_hub import hf_hub_download from safetensors.torch import load_file # 1. Classifier Architecture class Head(nn.Module): def __init__(self, c): super().__init__() self.pre_classifier, self.bn = nn.Linear(768, 768), nn.BatchNorm1d(768) self.classifier = nn.Linear(768, c) def forward(self, x): return self.classifier(torch.dropout(torch.relu(self.bn(self.pre_classifier(x))), 0.3, False)) # 2. Config REPO = "CrabInHoney/urlbert-tiny-v5" CLS_HEADS = { "MSMalicious-URLs-dataset_head.safetensors": {0: "BENIGN", 1: "DEFACEMENT", 2: "PHISHING", 3: "MALWARE"}, "cyPhishing-Email-Detection_head.safetensors": {0: "LEGIT EMAIL", 1: "PHISH EMAIL", 2: "LEGIT URL", 3: "PHISH URL"}, "PSSpam-Email-Classification_head.safetensors": {0: "LEGIT EMAIL", 1: "SPAM EMAIL"}, "kmPhishing-urls_head.safetensors": {0: "SAFE URL", 1: "PHISHING"}, "eaphishing-dataset_head.safetensors": {0: "SAFE", 1: "PHISHING"}, "zlphishing-email-dataset_head.safetensors": {0: "SAFE EMAIL", 1: "PHISH EMAIL"} } GEN_FILE = "heads/annotationGenHead.safetensors" # 3. Load Models print("Loading models...") tok = AutoTokenizer.from_pretrained(REPO) enc = AutoModel.from_pretrained(REPO) # Load Classifiers models = {} for f, lbls in CLS_HEADS.items(): h = Head(len(lbls)) h.load_state_dict(load_file(hf_hub_download(REPO, f"heads/{f}"))) h.eval() models[f] = (h, lbls) # Load Generator dec_conf = BertConfig(vocab_size=tok.vocab_size, hidden_size=256, num_hidden_layers=4, num_attention_heads=4, intermediate_size=1024, is_decoder=True, add_cross_attention=True) gen_model = EncoderDecoderModel(encoder=enc, decoder=AutoModelForCausalLM.from_config(dec_conf)) gen_model.load_state_dict(load_file(hf_hub_download(REPO, GEN_FILE)), strict=False) gen_model.eval() # 4. Inference text = "http://paypal-secure-login.update.com" inputs = tok(text, return_tensors="pt", truncation=True, max_length=512) print(f"Target: {text}\n") print(f"{'HEAD':<30} {'VERDICT':<15} {'CONF'}") with torch.no_grad(): # Run Classifiers emb = enc(**inputs).last_hidden_state[:, 0, :] for fname, (model, labels) in models.items(): probs = F.softmax(model(emb), dim=1)[0] top_id = probs.argmax().item() verdict = labels[top_id] c = "\033[91m" if any(x in verdict for x in ["PHISH", "MALWARE", "SPAM", "DEFACE"]) else "\033[92m" print(f"{fname.split('_')[0]:<30} {c}{verdict:<15}\033[0m {probs[top_id]:.1%}") # Run Generator (Fix: Explicitly pass decoder_start_token_id) print("-" * 55) out = gen_model.generate( inputs.input_ids, max_length=60, num_beams=5, decoder_start_token_id=tok.cls_token_id, eos_token_id=tok.sep_token_id ) desc = tok.decode(out[0], skip_special_tokens=True) print(f"Generated Description: \033[96m{desc}\033[0m") ``` Output: ``` HEAD VERDICT CONF MSMalicious-URLs-dataset PHISHING 100.0% cyPhishing-Email-Detection PHISH URL 99.1% PSSpam-Email-Classification SPAM EMAIL 99.9% kmPhishing-urls PHISHING 91.8% eaphishing-dataset PHISHING 100.0% zlphishing-email-dataset PHISH EMAIL 55.2% ------------------------------------------------------- Generated Description: financial institution phishing ```