Phishing Email Detection Distilbert V2.1

🚀 DistilBERTベースのフィッシングメール検出モデル

このモデルは、メールやURLが安全か潜在的なフィッシングかを多ラベル分類するために、DistilBERTをベースに微調整されています。

🚀 クイックスタート

このモデルはDistilBERTをベースに構築され、メールとURLの多ラベル分類に微調整されています。以下に、使い方のサンプルコードを示します。

from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch

# モデルとトークナイザーをロード
tokenizer = AutoTokenizer.from_pretrained("your-username/model-name")
model = AutoModelForSequenceClassification.from_pretrained("your-username/model-name")

def predict_email(email_text):
    # 前処理とトークナイズ
    inputs = tokenizer(
        email_text,
        return_tensors="pt",
        truncation=True,
        max_length=512
    )
    
    # 予測を取得
    with torch.no_grad():
        outputs = model(**inputs)
        predictions = torch.nn.functional.softmax(outputs.logits, dim=-1)
    
    # 各クラスの確率を取得
    probs = predictions[0].tolist()
    
    # ラベル辞書を作成
    labels = {
        "legitimate_email": probs[0],
        "phishing_url": probs[1],
        "legitimate_url": probs[2],
        "phishing_url_alt": probs[3]
    }
    
    # 最も可能性の高い分類を決定
    max_label = max(labels.items(), key=lambda x: x[1])
    
    return {
        "prediction": max_label[0],
        "confidence": max_label[1],
        "all_probabilities": labels
    }

✨ 主な機能

• メールとURLの多ラベル分類を行い、安全か潜在的なフィッシングかを判定します。
• DistilBERTをベースにしており、効率的な推論が可能です。

📦 インストール

pip install transformers
pip install torch

💻 使用例

基本的な使用法

# サンプルの使用例
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/PhishingEmailDetectionで入手可能です。

📄 ライセンス

このプロジェクトはApache-2.0ライセンスの下で提供されています。