user ──assumes_role──► role ──grants_access──► cloud
user ──accesses──────► device ──has_vuln──────► cve
user ──connects_to───► ip ──connects_to───────► device
device ──lateral_movement──► device
device ──deployed_on──► cloud
cve ──exploits──────────────► device
Sample Output
Node: user_4456
Type: user
Probability: 0.7231
Risk: 🔴 HIGH
Top Edge: lateral_movement
Influence per relation:
lateral_movement +0.21031
has_vuln +0.14203
assumes_role +0.08821
accesses −0.03102
How to Load
import torch, json
from huggingface_hub import hf_hub_download
weights_path = hf_hub_download(repo_id="adarsh-aur/Heterogenous_GNN", filename="model_final.pt")
config_path = hf_hub_download(repo_id="adarsh-aur/Heterogenous_GNN", filename="config.json")
withopen(config_path) as f:
cfg = json.load(f)
model = HeteroRGCN(
feat_dims = cfg["feat_dims"],
node_types = cfg["node_types"],
edge_types = [tuple(et) for et in cfg["edge_types"]],
hidden_dim = cfg["hidden_dim"],
out_dim = cfg["out_dim"],
num_layers = cfg["num_layers"],
dropout = cfg["dropout"],
adapter_rank = cfg["adapter_rank"],
predict_types = cfg["predict_types"],
)
model.load_state_dict(torch.load(weights_path, map_location="cpu"))
model.eval()