Spaces:
Sleeping
Sleeping
File size: 2,634 Bytes
c6e6f10 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 | import json
import torch
import torch.nn as nn
import torch.optim as optim
from models.gnn import PhishingGNN_Model
from pipeline.graph_engine import TopologicalGraphEngine
from config import *
def run_training_pipeline():
print("[*] Launching Production Defender V5 Training Run...")
# 1. Try to load real data, fallback to mock if missing
try:
with open("network_telemetry.json", "r") as f:
dataset = json.load(f)
print(f"[+] Loaded {len(dataset)} real telemetry logs.")
except FileNotFoundError:
print("[-] network_telemetry.json not found. Falling back to simulated context.")
dataset = [
{'ip': '185.220.101.5', 'domain': 'shadow-dns-bypass.net', 'asn': 44050, 'is_malicious': 1.0},
{'ip': '193.56.28.14', 'asn': 57099, 'is_malicious': 1.0},
{'ip': '8.8.8.8', 'domain': 'dns.google', 'asn': 15169, 'is_malicious': 0.0}
]
# 2. Extract logs and structure target threat labels
engine = TopologicalGraphEngine()
x_dict, edge_index_dict = engine.extract_and_build(dataset)
labels_list = []
for ip_str in engine.ip_map.keys():
matching_logs = [log for log in dataset if log.get('ip') == ip_str]
label = matching_logs[0].get('is_malicious', 0.0) if matching_logs else 0.0
labels_list.append([label])
labels = torch.tensor(labels_list, dtype=torch.float32)
# 3. Model setup
in_channels_dict = {'ip': 16, 'domain': 32, 'asn': 8, 'cert': 16}
model = PhishingGNN_Model(
metadata=GRAPH_METADATA,
in_channels_dict=in_channels_dict,
hidden_channels=HIDDEN_CHANNELS,
num_heads=NUM_HEADS,
num_layers=NUM_LAYERS,
dropout_rate=DROPOUT_RATE
)
optimizer = optim.AdamW(model.parameters(), lr=0.0005, weight_decay=1e-3)
criterion = nn.BCELoss()
# 4. Optimization Loop
model.train()
for epoch in range(100):
optimizer.zero_grad()
predictions = model(x_dict, edge_index_dict)
# Only calculate loss on the nodes we have labels for
valid_preds = predictions[:len(labels)]
loss = criterion(valid_preds, labels)
loss.backward()
optimizer.step()
if (epoch + 1) % 20 == 0:
print(f"Epoch {epoch+1:03d}/100 | Topological Loss: {loss.item():.5f}")
model.safe_save(MODEL_SAVE_PATH)
print(f"[+] Operational weights successfully frozen at: {MODEL_SAVE_PATH}")
if __name__ == "__main__":
run_training_pipeline() |