Update app.py

app.py

@@ -112,11 +112,11 @@ def get_random_negative_node(source_node, destination_node, total_nodes):
 
    
 def predict(u_id, i_id, timestamp):
-    # Convert inputs to batch format
-    #u_tensor = torch.tensor([int(u_id)], dtype=torch.long)  # Single source node, batch size 1
-    #i_tensor = torch.tensor([int(i_id)], dtype=torch.long)  # Single destination node, batch size 1
-    #ts_tensor = torch.LongTensor([int(timestamp)])
+    # Before prediction
+    tgn.memory.__init_memory__()  # Re-initialize memory
+    tgn.embeddings.__init_memory__()
 
+    # Then run prediction
     u_array = np.array([int(u_id)])  # List of source nodes
     i_array = np.array([int(i_id)])  # List of destination nodes
     ts_array = np.array([int(timestamp)])  # List of timestamps
@@ -144,14 +144,8 @@ def predict(u_id, i_id, timestamp):
     # Negative node sampling: choose random negative nodes for this example (should be handled in your dataset)
     total_nodes = len(full_data.unique_nodes)  # or whatever gives your total node count
     random_negative_node = get_random_negative_node(u_id, i_id, total_nodes)
-    #negative_nodes_tensor = torch.tensor([random_negative_node], dtype=torch.long)
     negative_nodes_array = np.array([random_negative_node])  # List of negative nodes
 
-    # Make sure you batch them together properly (even if it is just a single edge for now)
-    #positive_probs, negative_probs = tgn.compute_edge_probabilities(
-    #        u_array, i_array, negative_nodes_array, ts_array, edge_idx_array   
-    #)
-
     # Call compute_edge_probabilities
     # You can pass edge_idxs or dummy edge features here depending on how your TGN expects it
     try:
@@ -183,18 +177,6 @@ def predict(u_id, i_id, timestamp):
     return f"Top {len(top_values)} predicted interaction probabilities:\n{formatted_result}"
 
 
-#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).",
-#)
-
 with gr.Blocks() as demo:
     gr.Markdown("## 🧠 TGN Playground (Wikipedia)")
     gr.Markdown("Enter node IDs and timestamp to predict future interaction probability using Temporal Graph Networks (TGN).")