urlbert-tiny-v5

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:

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

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))

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"

print("Loading models...")
tok = AutoTokenizer.from_pretrained(REPO)
enc = AutoModel.from_pretrained(REPO)

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)

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()

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%}")

    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