Update app.py

app.py

@@ -111,9 +111,10 @@ def get_random_negative_node(source_node, destination_node, total_nodes):
     return random.choice(candidates)
 
    
-def predict(u_id, i_id, timestamp):
+def predict(u_id, i_id, timestamp, mem):
     # Before prediction
-    tgn.memory.__init_memory__()  # Re-initialize memory
+    if mem:
+        tgn.memory.__init_memory__()  # Re-initialize memory
 
     # Then run prediction
     u_array = np.array([int(u_id)])  # List of source nodes
@@ -183,12 +184,13 @@ with gr.Blocks() as demo:
     with gr.Row():
         src = gr.Number(label="Source Node ID")
         dst = gr.Number(label="Destination Node ID")
-        ts = gr.Number(label="Timestamp")
+        ts  = gr.Number(label="Timestamp")
+        mem = gr.Checkbox(label="Reinitialize memory, don't update on prediction!")  # 👈 your checkbox
 
     predict_btn = gr.Button("Predict")
     result = gr.Textbox(label="Prediction Output")
 
-    predict_btn.click(fn=predict, inputs=[src, dst, ts], outputs=result)
+    predict_btn.click(fn=predict, inputs=[src, dst, ts, mem], outputs=result)
 
     gr.Markdown("### Download Log File")
     log_file = gr.File(label="Log File", interactive=False)