created app.py

app.py

@@ -0,0 +1,54 @@
+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()