| import json |
| import math |
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| import gradio as gr |
|
|
| with open("config.json") as f: |
| CFG = json.load(f) |
|
|
| with open("tokenizer.json") as f: |
| tokenizer = json.load(f) |
|
|
| with open("threshold.txt") as f: |
| THRESHOLD = float(f.read().strip()) |
|
|
| CHAR2ID = tokenizer["char2id"] |
|
|
| class PositionalEncoding(nn.Module): |
| def __init__(self, d_model, max_len=512, dropout=0.1): |
| super().__init__() |
| self.dropout = nn.Dropout(dropout) |
|
|
| pe = torch.zeros(max_len, d_model) |
|
|
| position = torch.arange(0, max_len).unsqueeze(1) |
|
|
| div_term = torch.exp( |
| torch.arange(0, d_model, 2) |
| * (-math.log(10000.0) / d_model) |
| ) |
|
|
| pe[:, 0::2] = torch.sin(position * div_term) |
| pe[:, 1::2] = torch.cos(position * div_term) |
|
|
| self.register_buffer("pe", pe.unsqueeze(0)) |
|
|
| def forward(self, x): |
| return self.dropout(x + self.pe[:, :x.size(1)]) |
|
|
| class PhishingTransformer(nn.Module): |
|
|
| def __init__(self): |
| super().__init__() |
|
|
| self.embed = nn.Embedding( |
| CFG["vocab_size"], |
| CFG["embed_dim"], |
| padding_idx=CFG["pad_token"] |
| ) |
|
|
| self.pos = PositionalEncoding( |
| CFG["embed_dim"], |
| CFG["max_len"], |
| CFG["dropout"] |
| ) |
|
|
| layer = nn.TransformerEncoderLayer( |
| d_model=CFG["embed_dim"], |
| nhead=CFG["num_heads"], |
| dim_feedforward=CFG["ff_dim"], |
| dropout=CFG["dropout"], |
| batch_first=True, |
| activation="gelu", |
| norm_first=True |
| ) |
|
|
| self.encoder = nn.TransformerEncoder( |
| layer, |
| num_layers=CFG["num_layers"] |
| ) |
|
|
| self.norm = nn.LayerNorm(CFG["embed_dim"]) |
|
|
| self.head = nn.Sequential( |
| nn.Linear(CFG["embed_dim"], CFG["embed_dim"] // 2), |
| nn.GELU(), |
| nn.Dropout(CFG["dropout"]), |
| nn.Linear(CFG["embed_dim"] // 2, 2) |
| ) |
|
|
| def forward(self, x): |
|
|
| mask = (x == CFG["pad_token"]) |
|
|
| x = self.embed(x) |
|
|
| x = self.pos(x) |
|
|
| x = self.encoder( |
| x, |
| src_key_padding_mask=mask |
| ) |
|
|
| valid = (~mask).unsqueeze(-1).float() |
|
|
| x = (x * valid).sum(1) / valid.sum(1).clamp(min=1e-9) |
|
|
| x = self.norm(x) |
|
|
| return self.head(x) |
|
|
| model = PhishingTransformer() |
|
|
| model.load_state_dict( |
| torch.load( |
| "phishing_detector.pth", |
| map_location="cpu" |
| ) |
| ) |
|
|
| model.eval() |
|
|
| def predict(url): |
|
|
| ids = [ |
| CHAR2ID.get(c, 1) |
| for c in url[:CFG["max_len"]] |
| ] |
|
|
| while len(ids) < CFG["max_len"]: |
| ids.append(CFG["pad_token"]) |
|
|
| x = torch.tensor([ids]) |
|
|
| with torch.no_grad(): |
| logits = model(x) |
| probs = F.softmax(logits, dim=-1)[0] |
|
|
| score = float(probs[1]) |
|
|
| return { |
| "phishing": score >= THRESHOLD, |
| "score": score |
| } |
|
|
| demo = gr.Interface( |
| fn=predict, |
| inputs="text", |
| outputs="json", |
| title="Phishing Detector" |
| ) |
|
|
| demo.launch() |
