Phishing Email Detection Distilbert V2.4.1

🚀 基於DistilBERT的釣魚郵件檢測模型

本模型基於DistilBERT構建，可對郵件和URL進行多標籤分類，判斷其是否安全或存在釣魚風險，為網絡安全提供有力保障。

✨ 主要特性

• 基於DistilBERT架構，具備高效的特徵提取能力。
• 經過微調，適用於郵件和URL的多標籤分類任務。
• 利用Hugging Face Trainer API進行微調，訓練過程更加便捷。

📦 安裝指南

pip install transformers
pip install torch

💻 使用示例

基礎用法

from transformers import AutoTokenizer, AutoModelForSequenceClassification
tokenizer = AutoTokenizer.from_pretrained("cybersectony/phishing-email-detection-distilbert_v2.4.1")
import torch

# Load model and tokenizer
model = AutoModelForSequenceClassification.from_pretrained("cybersectony/phishing-email-detection-distilbert_v2.4.1")

def predict_email(email_text):
    # Preprocess and tokenize
    inputs = tokenizer(
        email_text,
        return_tensors="pt",
        truncation=True,
        max_length=512
    )
    
    # Get prediction
    with torch.no_grad():
        outputs = model(**inputs)
        predictions = torch.nn.functional.softmax(outputs.logits, dim=-1)
    
    # Get probabilities for each class
    probs = predictions[0].tolist()
    
    # Create labels dictionary
    labels = {
        "legitimate_email": probs[0],
        "phishing_url": probs[1],
        "legitimate_url": probs[2],
        "phishing_url_alt": probs[3]
    }
    
    # Determine the most likely classification
    max_label = max(labels.items(), key=lambda x: x[1])
    
    return {
        "prediction": max_label[0],
        "confidence": max_label[1],
        "all_probabilities": labels
    }

高級用法

# 示例用法
email = """
Dear User,
Your account security needs immediate attention. Please verify your credentials.
Click here: http://suspicious-link.com
"""

result = predict_email(email)
print(f"Prediction: {result['prediction']}")
print(f"Confidence: {result['confidence']:.2%}")
print("\nAll probabilities:")
for label, prob in result['all_probabilities'].items():
    print(f"{label}: {prob:.2%}")

📚 詳細文檔

模型概述

此模型基於DistilBERT，經過微調後可對郵件和URL進行多標籤分類，判斷其是否安全或存在釣魚風險。

關鍵規格

性能指標

數據集詳情

該模型在一個自定義的郵件和URL數據集上進行訓練，這些數據被標記為合法或釣魚。該數據集可在Hugging Face Hub上的cybersectony/PhishingEmailDetectionv2.0獲取。

📄 許可證

本項目採用Apache 2.0許可證。