Merge step 5 & 6

app.py

@@ -24,8 +24,7 @@ from backend import (
     pre_process_encrypt_send_bank,
     pre_process_encrypt_send_third_party,
     run_fhe,
-    get_output,
-    decrypt_output,
+    get_output_and_decrypt,
     years_employed_encrypt_run_decrypt,
 )
 
@@ -85,7 +84,7 @@ with demo:
             gr.Markdown("### User")
             bool_inputs = gr.CheckboxGroup(["Car", "Property", "Mobile phone"], label="Which of the following do you actively hold or own?")
             num_children = gr.Slider(**CHILDREN_MIN_MAX, step=1, label="Number of children", info="How many children do you have ?")
-            household_size = gr.Slider(**FAMILY_MIN_MAX, step=1, label="Household size", info="How many members does your household have? ?")
+            household_size = gr.Slider(**FAMILY_MIN_MAX, step=1, label="Household size", info="How many members does your household have ?")
             total_income = gr.Slider(**INCOME_MIN_MAX, label="Income", info="What's you total yearly income (in euros) ?")
             age = gr.Slider(**AGE_MIN_MAX, step=1, label="Age", info="How old are you ?")
             income_type = gr.Dropdown(choices=INCOME_TYPES, value=INCOME_TYPES[0], label="Income type", info="What is your main type of income ?")
@@ -184,7 +183,7 @@ with demo:
         """
     )
 
-    gr.Markdown("## Step 5: Receive the encrypted output from the server.")
+    gr.Markdown("## Step 5: Receive the encrypted output from the server and decrypt.")
     gr.Markdown(
         """
         The value displayed below is a shortened byte representation of the actual encrypted output.
@@ -195,34 +194,18 @@ with demo:
     encrypted_output_representation = gr.Textbox(
         label="Encrypted output representation: ", max_lines=2, interactive=False
     )
-
-    # Button to send the encodings to the server using post method
-    get_output_button.click(
-        get_output, 
-        inputs=[client_id], 
-        outputs=[encrypted_output_representation],
-    )
-
-    gr.Markdown("## Step 6: Decrypt the output.")
-    gr.Markdown(
-        """
-        The user is able to decrypt the prediction using its private key.
-        """
-    )
-
-    decrypt_button = gr.Button("Decrypt the output")
     prediction_output = gr.Textbox(
         label="Prediction", max_lines=1, interactive=False
     )
 
-    # Button to decrypt the output
-    decrypt_button.click(
-        decrypt_output,
-        inputs=[client_id],
-        outputs=[prediction_output],
+    # Button to send the encodings to the server using post method
+    get_output_button.click(
+        get_output_and_decrypt, 
+        inputs=[client_id], 
+        outputs=[prediction_output, encrypted_output_representation],
     )
 
-    gr.Markdown("## Step 7 (optional): Explain the prediction.")
+    gr.Markdown("## Step 6 (optional): Explain the prediction.")
     gr.Markdown(
         """
         In case the credit card is likely to be denied, the user can run a second model in order to 

backend.py

@@ -410,14 +410,15 @@ def run_fhe(client_id):
             raise gr.Error("Please send the inputs from all three parties to the server first.")
 
 
-def get_output(client_id):
+def get_output_and_decrypt(client_id):
     """Retrieve the encrypted output.
 
     Args:
         client_id (str): The current client ID to consider.
     
     Returns:
-        encrypted_output_short (bytes): A byte short representation of the encrypted output. 
+        (Tuple[str, bytes]): The output message based on the decrypted prediction as well as 
+            a byte short representation of the encrypted output. 
     """
 
     if client_id == "":
@@ -434,56 +435,28 @@ def get_output(client_id):
         data=data,
     ) as response:
         if response.ok:
-            encrypted_output = response.content
-
-            # Save the encrypted output to bytes in a file as it is too large to pass through regular
-            # Gradio buttons (see https://github.com/gradio-app/gradio/issues/1877)
-            encrypted_output_path = _get_client_file_path("encrypted_output", client_id)
-
-            with encrypted_output_path.open("wb") as encrypted_output_file:
-                encrypted_output_file.write(encrypted_output)
+            encrypted_output_proba = response.content
             
             # Create a truncated version of the encrypted inputs for display
-            encrypted_output_short = shorten_bytes_object(encrypted_output)
-
-            return encrypted_output_short
-        else:
-            raise gr.Error("Please run the FHE execution first and wait for it to be completed.")
-
+            encrypted_output_short = shorten_bytes_object(encrypted_output_proba)
 
-def decrypt_output(client_id):
-    """Decrypt the result.
+            # Retrieve the client API
+            client = _get_client(client_id)
 
-    Args:
-        client_id (str): The current client ID to consider.
-
-    Returns:
-        output(numpy.ndarray): The decrypted output
-    """
+            # Deserialize, decrypt and post-process the encrypted output
+            output_proba = client.deserialize_decrypt_dequantize(encrypted_output_proba)
 
-    if client_id == "":
-        raise gr.Error("Please generate the keys first.")
-
-    # Get the encrypted output path
-    encrypted_output_path = _get_client_file_path("encrypted_output", client_id)
-
-    if not encrypted_output_path.is_file():
-        raise gr.Error("Please receive the outputs from the server first.")
-
-    # Load the encrypted output as bytes
-    with encrypted_output_path.open("rb") as encrypted_output_file:
-        encrypted_output_proba = encrypted_output_file.read()
-
-    # Retrieve the client API
-    client = _get_client(client_id)
-
-    # Deserialize, decrypt and post-process the encrypted output
-    output_proba = client.deserialize_decrypt_dequantize(encrypted_output_proba)
+            # Determine the predicted class
+            output = numpy.argmax(output_proba, axis=1).squeeze()
+            
+            return (
+                "Credit card is likely to be approved ✅" if output == 1 
+                else "Credit card is likely to be denied ❌",
+                encrypted_output_short,
+            )
 
-    # Determine the predicted class
-    output = numpy.argmax(output_proba, axis=1).squeeze()
-    
-    return "Credit card is likely to be approved ✅" if output == 1 else "Credit card is likely to be denied ❌"
+        else:
+            raise gr.Error("Please run the FHE execution first and wait for it to be completed.")
 
 
 def years_employed_encrypt_run_decrypt(client_id, prediction_output, *inputs):