import spaces # MUST come before torch / any CUDA-touching import import os import torch import torch.nn as nn import numpy as np import gradio as gr from transformers import PreTrainedTokenizerFast # --------------------------------------------------------------------------- # Model definition (ported from the official DRIFT repository) # https://github.com/snsec-net/2026-DSN-DRIFT · model.py # --------------------------------------------------------------------------- class TokenEmbedding(nn.Module): def __init__(self, vocab_size, d_model, padding_idx): super().__init__() self.embedding = nn.Embedding(vocab_size, d_model, padding_idx=padding_idx) torch.nn.init.xavier_normal_(self.embedding.weight) def forward(self, input): return self.embedding(input) class PositionalEncoding(nn.Module): def __init__(self, d_model, max_len): super().__init__() self.pos_embed = nn.Embedding(max_len, d_model) torch.nn.init.xavier_normal_(self.pos_embed.weight) def forward(self, x): B, L, _ = x.size() device = x.device pos_ids = torch.arange(L, device=device).unsqueeze(0).expand(B, L) return x + self.pos_embed(pos_ids) class Transformer(nn.Module): def __init__(self, d_model, n_heads, dim_feedforward, num_layers, dropout=0.1): super().__init__() encoder_layer = nn.TransformerEncoderLayer( d_model=d_model, nhead=n_heads, dim_feedforward=dim_feedforward, dropout=dropout, batch_first=True, ) self.encoder = nn.TransformerEncoder(encoder_layer, num_layers) self.dropout = nn.Dropout(dropout) def forward(self, x, mask=None): if mask is not None and mask.any(): out = self.encoder(x, src_key_padding_mask=mask) else: out = self.encoder(x) return self.dropout(out) class PretrainedModel(nn.Module): def __init__(self, vocab_size, d_model, n_heads, dim_feedforward, num_layers, max_len, dropout=0.1, padding_idx=0, tov_norm='pool'): super().__init__() self.d_model = d_model self.padding_idx = padding_idx self.max_len = max_len self.tov_norm = tov_norm self.embedding = TokenEmbedding(vocab_size, d_model, padding_idx) self.positional_encoding = PositionalEncoding(d_model, max_len) self.transformer = Transformer(d_model, n_heads, dim_feedforward, num_layers, dropout) def create_padding_mask(self, input_ids): return (input_ids == self.padding_idx) class FinetuningHead(nn.Module): def __init__(self, input_dim, d_model, dropout): super().__init__() self.input_dim = input_dim self.d_model = d_model self.dense1 = nn.Linear(input_dim, d_model * 2) self.dropout = nn.Dropout(dropout) self.classifier = nn.Linear(d_model * 2, 2) def forward(self, encoder_output): x = self.dropout(encoder_output) x = self.dense1(x) x = torch.relu(x) x = self.dropout(x) logits = self.classifier(x) return logits class FineTuningModel(nn.Module): def __init__(self, pretrain_model_t=None, pretrain_model_c=None, dropout=0.1, padding_idx=0, clf_norm='pool', freeze_backbone=False): super().__init__() self.padding_idx = padding_idx self.clf_norm = clf_norm self.use_token = pretrain_model_t is not None self.use_char = pretrain_model_c is not None sample_model = pretrain_model_t if self.use_token else pretrain_model_c d_model = sample_model.d_model num_active_paths = sum([self.use_token, self.use_char]) dim_per_path = d_model * 2 if clf_norm == 'pool' else d_model total_input_dim = dim_per_path * num_active_paths if self.use_token: self.transformer_encoder_t = pretrain_model_t.transformer self.embedding_t = pretrain_model_t.embedding self.positional_encoding_t = pretrain_model_t.positional_encoding if self.use_char: self.transformer_encoder_c = pretrain_model_c.transformer self.embedding_c = pretrain_model_c.embedding self.positional_encoding_c = pretrain_model_c.positional_encoding self.classifier_head = FinetuningHead( input_dim=total_input_dim, d_model=d_model, dropout=dropout ) def create_padding_mask(self, input_ids): return (input_ids == self.padding_idx).to(input_ids.device) def forward(self, input_ids_t=None, input_ids_c=None): features = [] if self.use_token and input_ids_t is not None: t_embed = self.embedding_t(input_ids_t) t_x = self.positional_encoding_t(t_embed) t_mask = self.create_padding_mask(input_ids_t) t_out = self.transformer_encoder_t(t_x, mask=t_mask) if self.clf_norm == 'pool': valid_mask = (~t_mask).float().unsqueeze(-1) t_sum = (t_out * valid_mask).sum(dim=1) t_len = valid_mask.sum(dim=1).clamp(min=1) t_mean = t_sum / t_len t_max = (t_out.masked_fill(valid_mask == 0, -1e9)).max(dim=1).values t_feat = torch.cat([t_max, t_mean], dim=1) else: t_feat = t_out[:, 0, :] features.append(t_feat) if self.use_char and input_ids_c is not None: c_embed = self.embedding_c(input_ids_c) c_x = self.positional_encoding_c(c_embed) c_mask = self.create_padding_mask(input_ids_c) c_out = self.transformer_encoder_c(c_x, mask=c_mask) if self.clf_norm == 'pool': valid_mask = (~c_mask).float().unsqueeze(-1) c_sum = (c_out * valid_mask).sum(dim=1) c_len = valid_mask.sum(dim=1).clamp(min=1) c_mean = c_sum / c_len c_max = (c_out.masked_fill(valid_mask == 0, -1e9)).max(dim=1).values c_feat = torch.cat([c_max, c_mean], dim=1) else: c_feat = c_out[:, 0, :] features.append(c_feat) combined_output = torch.cat(features, dim=1) if len(features) > 1 else features[0] return self.classifier_head(combined_output) # --------------------------------------------------------------------------- # Configuration constants (from utility/config.py) # --------------------------------------------------------------------------- D_MODEL = 256 N_HEADS = 8 NUM_LAYERS = 12 DIM_FEEDFORWARD = 768 MAX_LEN_SUBWORD = 30 MAX_LEN_CHAR = 77 VOCAB_SIZE_SUBWORD = 30522 VOCAB_SIZE_CHAR = 43 PADDING_IDX = 0 CLF_NORM = 'pool' # Special token IDs (from preprocessing.py SpecialIDs) PAD_ID = 0 UNK_ID = 1 CLS_ID = 2 SEP_ID = 3 MASK_ID = 4 CHAR_LIST = list("abcdefghijklmnopqrstuvwxyz0123456789-.") SPECIAL_TOKENS = ['[PAD]', '[UNK]', '[CLS]', '[SEP]', '[MASK]'] ALL_TOKENS = SPECIAL_TOKENS + CHAR_LIST CHAR2ID = {char: idx for idx, char in enumerate(ALL_TOKENS)} MODEL_ID = "snsec-net/dga-detector-drift26dsn" # --------------------------------------------------------------------------- # Load tokenizer and model at module scope # --------------------------------------------------------------------------- TOKENIZER_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "tokenizer-0-30522-both.json") tokenizer = PreTrainedTokenizerFast(tokenizer_file=TOKENIZER_PATH) pt_model_c = PretrainedModel( vocab_size=VOCAB_SIZE_CHAR, d_model=D_MODEL, n_heads=N_HEADS, dim_feedforward=DIM_FEEDFORWARD, num_layers=NUM_LAYERS, max_len=MAX_LEN_CHAR, ) pt_model_t = PretrainedModel( vocab_size=VOCAB_SIZE_SUBWORD, d_model=D_MODEL, n_heads=N_HEADS, dim_feedforward=DIM_FEEDFORWARD, num_layers=NUM_LAYERS, max_len=MAX_LEN_SUBWORD, ) model = FineTuningModel(pt_model_t, pt_model_c, clf_norm=CLF_NORM) CKPT_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "finetuning.pt") state = torch.load(CKPT_PATH, map_location="cpu", weights_only=False) model.load_state_dict(state, strict=False) model = model.to("cuda").eval() # --------------------------------------------------------------------------- # Preprocessing (ported from preprocessing.py FineTuningDataset) # --------------------------------------------------------------------------- def domain_to_char_ids(domain: str) -> np.ndarray: """Character-level encoding: [CLS] + chars + [SEP], padded to MAX_LEN_CHAR.""" domain = domain.lower() token_indices = [CHAR2ID.get(c, UNK_ID) for c in domain] if len(token_indices) > MAX_LEN_CHAR - 2: token_indices = token_indices[:MAX_LEN_CHAR - 2] ids = [CLS_ID] + token_indices + [SEP_ID] if len(ids) < MAX_LEN_CHAR: ids += [PAD_ID] * (MAX_LEN_CHAR - len(ids)) return np.array(ids, dtype=np.int64) def domain_to_subword_ids(domain: str) -> np.ndarray: """Subword-level encoding: [CLS] + subwords + [SEP], padded to MAX_LEN_SUBWORD.""" domain = domain.lower() encoded = tokenizer(domain, add_special_tokens=False) token_indices = encoded["input_ids"] if len(token_indices) > MAX_LEN_SUBWORD - 2: token_indices = token_indices[:MAX_LEN_SUBWORD - 2] ids = [CLS_ID] + token_indices + [SEP_ID] if len(ids) < MAX_LEN_SUBWORD: ids += [PAD_ID] * (MAX_LEN_SUBWORD - len(ids)) return np.array(ids, dtype=np.int64) # --------------------------------------------------------------------------- # Inference # --------------------------------------------------------------------------- @spaces.GPU(duration=30) def detect_dga(domain: str): """Classify a domain name as Benign or DGA-generated. Args: domain: A domain name (e.g. "google.com" or "xkqjhfwiwbfw.info"). Returns: A tuple of (label_dict, confidence, benign_prob, dga_prob). """ domain = domain.strip() if not domain: return {"—": 1.0}, 0.0, 0.0, 0.0 char_ids = domain_to_char_ids(domain) subword_ids = domain_to_subword_ids(domain) x_t = torch.tensor(np.array([subword_ids]), dtype=torch.long, device="cuda") x_c = torch.tensor(np.array([char_ids]), dtype=torch.long, device="cuda") with torch.no_grad(): logits = model(x_t, x_c) probs = torch.softmax(logits, dim=1) pred = torch.argmax(logits, dim=1).item() benign_prob = probs[0, 0].item() dga_prob = probs[0, 1].item() label = "DGA" if pred == 1 else "Benign" confidence = max(benign_prob, dga_prob) return {label: confidence}, confidence, benign_prob, dga_prob # --------------------------------------------------------------------------- # Gradio UI # --------------------------------------------------------------------------- CSS = """ #col-container { max-width: 900px; margin: 0 auto; } .dark .gradio-container { color: var(--body-text-color); } """ with gr.Blocks() as demo: gr.Markdown( "# DRIFT: Drift-Resilient Invariant-Feature Transformer for DGA Detection\n" "Enter a domain name (effective second-level domain) to classify it as " "**Benign** or **DGA-generated**. DRIFT uses a dual-branch Transformer " "(character + subword) to learn invariant structural features that remain " "robust against concept drift.\n\n" "[Paper](https://huggingface.co/papers/2605.10436) · " "[GitHub](https://github.com/snsec-net/2026-DSN-DRIFT) · " "[Model](https://huggingface.co/snsec-net/dga-detector-drift26dsn)" ) with gr.Column(elem_id="col-container"): domain_input = gr.Textbox( label="Domain name (effective second-level domain)", placeholder="e.g. google or xkqjhfwiwbfw", info="Enter the effective second-level domain (eSLD) — TLD stripped, " "lowercased, characters a–z, 0–9, '-', '.' only.", ) run_btn = gr.Button("Classify", variant="primary") with gr.Row(): label_out = gr.Label(label="Classification") confidence_out = gr.Number(label="Confidence", precision=4) with gr.Accordion("Class probabilities", open=False): benign_bar = gr.Number(label="P(Benign)", precision=4) dga_bar = gr.Number(label="P(DGA)", precision=4) gr.Examples( examples=[ ["google"], ["github"], ["xkqjhfwiwbfw"], ["my-secure-login-portal"], ["qpalzm-xnkvjf"], ], inputs=[domain_input], fn=detect_dga, outputs=[label_out, confidence_out, benign_bar, dga_bar], cache_examples=True, cache_mode="lazy", ) run_btn.click( fn=detect_dga, inputs=[domain_input], outputs=[label_out, confidence_out, benign_bar, dga_bar], api_name="detect_dga", ) if __name__ == "__main__": demo.launch(mcp_server=True, theme=gr.themes.Citrus(), css=CSS)