Update predictor.py (#3)

- Update predictor.py (5a3b7b14551571ba0955b56a644f9968440c30e8)


Co-authored-by: Sandro Ferroni <[email protected]>

predictor.py

@@ -3,9 +3,13 @@ import joblib
 import numpy as np
 from concrete.ml.deployment import FHEModelClient, FHEModelServer
 import logging
+import gradio as gr
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)
+key_already_generated_condition = False
+encrypted_data = None
+encrypted_prediction = None
 
 # Paths to required files
 SCALER_PATH = os.path.join("models", "scaler.pkl")
@@ -32,40 +36,90 @@ except FileNotFoundError:
 # Load evaluation keys
 evaluation_keys = client.get_serialized_evaluation_keys()
 
-def predict(input_data):
+def predict():
     """
     Perform a local prediction using the compiled FHE model.
 
-    Args:
-        input_data (dict): User input data as a dictionary.
-
     Returns:
-        str: Prediction result ("Fraudulent" or "Non-fraudulent").
+        str: The prediction result.
+        str: A message indicating the status of the prediction.
     """
+    global encrypted_data, encrypted_prediction
+    if encrypted_data is None:
+        return None, gr.update(value="No encrypted data to predict. Please provide encrypted data ❌")
     try:
-        logging.info(f"Input Data: {input_data}")
-
-        # Scale the input data
-        scaled_data = scaler.transform([list(input_data.values())])
-        logging.info(f"Scaled Data: {scaled_data}")
-
-        # Encrypt the scaled data
-        encrypted_data = client.quantize_encrypt_serialize(scaled_data)
-        logging.info("Data encrypted successfully.")
-
         # Execute the model locally on encrypted data
         encrypted_prediction = server.run(
             encrypted_data, serialized_evaluation_keys=evaluation_keys
         )
         logging.info(f"Encrypted Prediction: {encrypted_prediction}")
+        return encrypted_prediction, gr.update(value="FHE evaluation is done. ✅")
+        
+    except Exception as e:
+        logging.error(f"Error during prediction: {e}")
+        return None, gr.update(value="No encrypted data to predict. Please provide encrypted data ❌")
+    
+def decrypt_prediction():
+    """
+    Decrypt and interpret the prediction result.
 
+    Returns:
+        str: The interpreted prediction result.
+    """
+    global encrypted_prediction
+    if encrypted_prediction is None:
+        return "No prediction to decrypt. Please make a prediction first. ❌", "No prediction to decrypt. Please make a prediction first. ❌"
+    try:
         # Decrypt the prediction result
         decrypted_prediction = client.deserialize_decrypt_dequantize(encrypted_prediction)
         logging.info(f"Decrypted Prediction: {decrypted_prediction}")
 
         # Interpret the prediction
         binary_prediction = int(np.argmax(decrypted_prediction))
-        return "Fraudulent" if binary_prediction == 1 else "Non-fraudulent"
+        return "⚠️ Fraudulent ⚠️" if binary_prediction == 1 else "😊 Non-fraudulent 😊", gr.update(value="Decryption successful ✅")
+        
     except Exception as e:
         logging.error(f"Error during prediction: {e}")
-        return "Error during prediction"
+        return "Error during prediction❌", "Error during prediction❌"
+
+def key_already_generated():
+    """
+    Check if the evaluation keys have already been generated.
+
+    Returns:
+        bool: True if the evaluation keys have already been generated, False otherwise.
+    """
+    global key_already_generated_condition
+    if evaluation_keys:
+        key_already_generated_condition = True
+        return True
+    return False
+
+def pre_process_encrypt_send_purchase(*inputs):
+    """
+    Pre-processes, encrypts, and sends the purchase data for prediction.
+
+    Args:
+        *inputs: Variable number of input arguments.
+
+    Returns:
+        (str): A short representation of the encrypted input to send in hex. 
+    """
+    global key_already_generated_condition, encrypted_data
+    if key_already_generated_condition == False:
+        return None, gr.update(value="Generate your key before. ❌")
+    try:
+        key_already_generated_condition = True
+        logging.info(f"Input Data: {inputs}")
+
+        # Scale the input data
+        scaled_data = scaler.transform([list(inputs)])
+        logging.info(f"Scaled Data: {scaled_data}")
+
+        # Encrypt the scaled data
+        encrypted_data = client.quantize_encrypt_serialize(scaled_data)
+        logging.info("Data encrypted successfully.")
+        return encrypted_data, gr.update(value="Inputs are encrypted and sent to server. ✅")
+    except Exception as e:
+        logging.error(f"Error during pre-processing: {e}")
+        return "Error during pre-processing"