File size: 1,532 Bytes
94de56f | 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 | import gradio as gr
import numpy as np
import torch
from utils.data_processing import get_data
from modules.tgn import TGN
# Load data
node_features, edge_features, full_data, train_data, val_data, test_data, \
new_node_val_data, new_node_test_data = get_data('wikipedia')
# Initialize model (must match training args)
tgn = TGN(
raw_message_dim=edge_features.shape[1],
memory_dim=172,
message_dim=100,
memory_update_at_end=False,
embedding_module_type='graph_attention',
message_function='identity',
memory_updater_type='gru',
n_heads=2,
dropout=0.1,
n_layers=1,
n_neighbors=10,
aggregator_type='last',
edge_features=edge_features,
node_features=node_features,
device='cpu',
use_memory=True,
)
tgn.load_state_dict(torch.load("tgn-attn-wikipedia.pth", map_location="cpu"))
tgn.eval()
def predict(u_id, i_id, timestamp):
u = np.array([int(u_id)])
i = np.array([int(i_id)])
ts = np.array([float(timestamp)])
prob = tgn.predict_edge_probabilities(u, i, ts)
return f"Predicted interaction probability: {prob[0]:.4f}"
demo = gr.Interface(
fn=predict,
inputs=[
gr.Number(label="Source Node ID"),
gr.Number(label="Destination Node ID"),
gr.Number(label="Timestamp"),
],
outputs="text",
title="🧠 TGN Playground (Wikipedia)",
description="Enter node IDs and timestamp to predict future interaction probability using Temporal Graph Networks (TGN).",
)
if __name__ == "__main__":
demo.launch()
|