Upload model.py with huggingface_hub

model.py

@@ -0,0 +1,76 @@
+"""
+DGA-CNN: Character-level CNN for DGA detection.
+Architecture from Patton et al. (adapted), trained on 54 DGA families.
+"""
+import string
+import torch
+import torch.nn as nn
+
+CHARS = string.ascii_lowercase + string.digits + "-._"
+CHAR2IDX = {c: i + 1 for i, c in enumerate(CHARS)}  # 0 = padding
+VOCAB_SIZE = len(CHARS) + 1  # 40
+MAXLEN = 75
+
+
+def encode_domain(domain: str) -> list:
+    domain = str(domain).lower().strip()
+    encoded = [CHAR2IDX.get(c, 0) for c in domain[:MAXLEN]]
+    return encoded + [0] * (MAXLEN - len(encoded))
+
+
+class DGACNN(nn.Module):
+    def __init__(self, vocab_size=VOCAB_SIZE, embedding_dim=32, num_classes=2):
+        super().__init__()
+        self.embedding = nn.Embedding(vocab_size, embedding_dim, padding_idx=0)
+        self.conv1 = nn.Conv1d(embedding_dim, 64, kernel_size=3, padding=1)
+        self.relu = nn.ReLU()
+        self.pool = nn.MaxPool1d(2)
+        self.dropout = nn.Dropout(0.3)
+        self.fc = nn.Linear(64 * (MAXLEN // 2), num_classes)
+
+    def forward(self, x):
+        x = self.embedding(x).transpose(1, 2)
+        x = self.pool(self.relu(self.conv1(x)))
+        x = x.view(x.size(0), -1)
+        x = self.dropout(x)
+        return self.fc(x)
+
+
+def load_model(weights_path: str, device: str = None):
+    """Load trained model from a local weights path."""
+    if device is None:
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+    model = DGACNN()
+    model.load_state_dict(torch.load(weights_path, map_location=device))
+    model.to(device)
+    model.eval()
+    return model
+
+
+def predict(model, domains, device: str = None, batch_size: int = 256):
+    """
+    Predict DGA vs legit for a list of domain strings.
+    Returns list of dicts: [{"domain": ..., "label": "dga"/"legit", "score": float}]
+    """
+    if device is None:
+        device = next(model.parameters()).device
+    if isinstance(domains, str):
+        domains = [domains]
+
+    results = []
+    for i in range(0, len(domains), batch_size):
+        batch = domains[i : i + batch_size]
+        encoded = [encode_domain(d) for d in batch]
+        x = torch.tensor(encoded, dtype=torch.long).to(device)
+        with torch.no_grad():
+            logits = model(x)
+            probs = torch.softmax(logits, dim=1)
+            preds = logits.argmax(dim=1).cpu().tolist()
+            scores = probs[:, 1].cpu().tolist()
+        for domain, pred, score in zip(batch, preds, scores):
+            results.append({
+                "domain": domain,
+                "label": "dga" if pred == 1 else "legit",
+                "score": round(score, 4),
+            })
+    return results