Upload app.py

app.py

@@ -70,7 +70,7 @@ def keygen():
     print("Initializing FHEModelClient...")
     # Let's create a user_id
     user_id = numpy.random.randint(0, 2**32)
-    fhe_api = FHEModelClient(f"deployment/deployment_{task}", f".fhe_keys/{user_id}")
+    fhe_api = FHEModelClient(f"fhe_model", f".fhe_keys/{user_id}")
     fhe_api.load()
 
     # Generate a fresh key
@@ -169,30 +169,35 @@ def decrypt_prediction(user_id):
     predictions = fhe_api.deserialize_decrypt_dequantize(encrypted_prediction)
 
 
-def update(name):
-    return f"Welcome to Gradio, {name}!"
+if __name__ == "__main__":
+    import gradio as gr
+
+
+def process_pipeline(task, test_file, user_id):
+    eval_key = keygen(task, user_id)
+    encodings = encode_quantize_encrypt(test_file, eval_key)
+    encrypted_quantized_encoding = run_fhe(encodings)
+    encrypted_prediction = decrypt_prediction(encrypted_quantized_encoding)
 
+    return eval_key, encodings, encrypted_quantized_encoding, encrypted_prediction
 
 if __name__ == "__main__":
+    import gradio as gr
     app = gr.Interface(
-        [
-            keygen,
-            encode_quantize_encrypt,
-            run_fhe,
-            decrypt_prediction,
-        ],
-        [
-            gr.inputs.Textbox(label="Task", default="malware"),
-            gr.inputs.File(label="Test File"),
-            gr.inputs.Textbox(label="User ID"),
+        fn=process_pipeline,
+        inputs=[
+            gr.Textbox(label="Task", value="malware"),
+            gr.File(label="Test File"),
+            gr.Textbox(label="User ID"),
         ],
-        [
-            gr.outputs.Textbox(label="Evaluation Key"),
-            gr.outputs.Textbox(label="Encodings"),
-            gr.outputs.Textbox(label="Encrypted Quantized Encoding"),
-            gr.outputs.Textbox(label="Encrypted Prediction"),
+        outputs=[
+            gr.Textbox(label="Evaluation Key"),
+            gr.Textbox(label="Encodings"),
+            gr.Textbox(label="Encrypted Quantized Encoding"),
+            gr.Textbox(label="Encrypted Prediction"),
         ],
         title="FHE Model",
         description="This is a FHE Model",
     )
-    app.launch()
+
+    app.launch()