license: apache-2.0 datasets:

• cybersectony/PhishingEmailDetectionv2.0 language:
• en base_model:
• distilbert/distilbert-base-uncased library_name: transformers

基于distilBERT的钓鱼邮件检测模型

模型概述

该模型基于DistilBERT架构，专为多标签分类任务设计，用于判断邮件和URL是否安全或存在钓鱼风险。

核心参数

• 基础架构: DistilBERT
• 任务类型: 多标签分类
• 微调框架: Hugging Face Trainer API
• 训练周期: 3轮

性能指标

• 准确率: 99.58
• F1分数: 99.579
• 精确率: 99.583
• 召回率: 99.58

数据集说明

模型训练采用定制数据集，包含标记为合法或钓鱼的邮件和URL样本。数据集存放于Hugging Face Hub的cybersectony/PhishingEmailDetectionv2.0。

使用指南

安装依赖

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

# 加载模型和分词器
model = AutoModelForSequenceClassification.from_pretrained("cybersectony/phishing-email-detection-distilbert_v2.4.1")

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 = {
        "合法邮件": probs[0],
        "钓鱼链接": probs[1],
        "合法链接": probs[2],
        "替代钓鱼链接": probs[3]
    }
    
    # 确定最高概率分类
    max_label = max(labels.items(), key=lambda x: x[1])
    
    return {
        "预测结果": max_label[0],
        "置信度": max_label[1],
        "全部分类概率": labels
    }

使用示例

# 示例用法
email = """
尊敬的客户：
您的账户安全需要立即关注，请验证您的凭据。
点击链接：http://可疑网站.com
"""

result = predict_email(email)
print(f"预测结果: {result['预测结果']}")
print(f"置信度: {result['置信度']:.2%}")
print("\n全部分类概率:")
for label, prob in result['全部分类概率'].items():
    print(f"{label}: {prob:.2%}")