Update app.py

app.py

@@ -128,7 +128,8 @@ def encode_quantize(test_file, eval_key, encodings):
 
     return encodings
 
-def encrypt_encoded_quantize(encodings, user_id, eval_key):
+def encrypt_encoded_quantize(encodings, eval_key):
+    global user_id
     fhe_api = FHEModelClient(f"fhe_model", f".fhe_keys/{user_id}")
     fhe_api.load()
     
@@ -146,7 +147,8 @@ def encrypt_encoded_quantize(encodings, user_id, eval_key):
     encrypted_quantized_encoding_shorten_hex = "".join(f"{i:02x}" for i in encrypted_quantized_encoding_shorten)
     return (encodings,quantized_encodings,encrypted_quantized_encoding_shorten_hex)
 
-def run_fhe(user_id):
+def run_fhe():
+    global user_id
     encoded_data_path = Path(f"tmp/tmp_encrypted_quantized_encoding_{user_id}.npy")
     encrypted_quantized_encoding = numpy.load(encoded_data_path)
 
@@ -178,7 +180,8 @@ def run_fhe(user_id):
     return encrypted_prediction_shorten_hex
 
 
-def decrypt_prediction(user_id):
+def decrypt_prediction():
+    global user_id
     encoded_data_path = Path(f"tmp/tmp_encrypted_prediction_{user_id}.npy")
 
     # Read encrypted_prediction from the file
@@ -348,7 +351,6 @@ if __name__ == "__main__":
             "The encrypted scan result is sent back to the client, who can finally decrypt it with their private key. Only the client is aware of the original file content and the scan result."
         )
         b_decrypt_result = gr.Button("🔍 Decrypt scan result")
-        user_id_input = gr.Number(value=user_id, visible=False)
         scan_result = gr.Textbox(label="Scan Result:")
         
         eval_key_input = gr.Textbox(value=eval_key, visible=False)
@@ -366,15 +368,15 @@ if __name__ == "__main__":
         # Button to encrypt file
         b_encrypt_file.click(
             fn=encrypt_encoded_quantize,
-            inputs=[extracted_vector, user_id_input, eval_key_input],
+            inputs=[extracted_vector, eval_key_input],
             outputs=[encrypted_file],
         )
 
         # Button to run FHE-based malware scan
-        b_run_fhe_scan.click(fn=run_fhe, inputs=[user_id_input], outputs=[encrypted_scan_result])
+        b_run_fhe_scan.click(fn=run_fhe, inputs=[], outputs=[encrypted_scan_result])
 
         # Button to decrypt the scan result
-        b_decrypt_result.click(fn=decrypt_prediction, inputs=[user_id_input], outputs=[scan_result])
+        b_decrypt_result.click(fn=decrypt_prediction, inputs=[], outputs=[scan_result])
 
         gr.Markdown(
             "ClairVault is built using advanced Fully Homomorphic Encryption techniques to ensure your data remains private and secure throughout the entire malware scanning process."