dviro/urlbert-onnx-mini

URLBert Phishing Detector (ONNX)

This is an ONNX version of the CrabInHoney/urlbert-tiny-v4-phishing-classifier model, which is designed to detect phishing URLs.

Usage with Hugging Face Transformers

from transformers import AutoTokenizer, AutoModelForSequenceClassification
import onnxruntime as ort
import numpy as np

# Load tokenizer
tokenizer = AutoTokenizer.from_pretrained("dviro/urlbert-onnx-mini")

# Load ONNX model
ort_session = ort.InferenceSession("phishing_detector.onnx")

# Example URL
url = "example.com"

# Tokenize
inputs = tokenizer(url, return_tensors="np", padding="max_length", truncation=True)
input_ids = inputs["input_ids"].astype(np.int64)
attention_mask = inputs["attention_mask"].astype(np.int64)
token_type_ids = inputs.get("token_type_ids", np.zeros_like(input_ids)).astype(np.int64)

# Run inference
outputs = ort_session.run(
    None, 
    {
        "input_ids": input_ids,
        "attention_mask": attention_mask,
        "token_type_ids": token_type_ids
    }
)

# Process results
logits = outputs[0][0]
scores = np.exp(logits) / np.sum(np.exp(logits))
labels = ["Safe", "Phishing"]
prediction = labels[np.argmax(scores)]
confidence = np.max(scores)

print(f"URL: {url}")
print(f"Prediction: {prediction}")
print(f"Confidence: {confidence:.4f}")

Usage with ONNX Runtime Web

This model can also be used in browser applications with ONNX Runtime Web. See the example code in the accompanying GitHub repository.

Model Details

• Original model: CrabInHoney/urlbert-tiny-v4-phishing-classifier
• Task: Binary classification (phishing detection)
• Framework: ONNX
• Model size: Small (BERT-tiny architecture)