Create predictor.py

predictor.py

@@ -0,0 +1,135 @@
+import os
+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")
+FHE_FILES_PATH = os.path.join("models", "fhe_files")
+
+# Load the scaler
+try:
+    scaler = joblib.load(SCALER_PATH)
+    logging.info("Scaler loaded successfully.")
+except FileNotFoundError:
+    logging.error(f"Error: The file scaler.pkl is missing at {SCALER_PATH}.")
+    raise
+
+# Initialize the FHE client and server
+try:
+    client = FHEModelClient(path_dir=FHE_FILES_PATH, key_dir=FHE_FILES_PATH)
+    server = FHEModelServer(path_dir=FHE_FILES_PATH)
+    server.load()
+    logging.info("FHE Client and Server initialized successfully.")
+except FileNotFoundError:
+    logging.error(f"Error: The FHE files (client.zip, server.zip) are missing in {FHE_FILES_PATH}.")
+    raise
+
+# Load evaluation keys
+evaluation_keys = client.get_serialized_evaluation_keys()
+
+def predict():
+    """
+    Perform a local prediction using the compiled FHE model.
+    Returns:
+        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:
+        # 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.hex(), 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))
+        # Ensure the prediction is a flat array
+        if isinstance(decrypted_prediction, np.ndarray) and decrypted_prediction.ndim > 1:
+            decrypted_prediction = decrypted_prediction.flatten()
+
+        # Generate the HTML for the percentage bar
+        bar_html = f"""
+        <div style="width: 100%; background-color: lightgray; border-radius: 5px; overflow: hidden; display: flex;">
+            <div style="width: {decrypted_prediction[0] * 100}%; background-color: green; color: white; text-align: center; padding: 5px 0;">
+                {decrypted_prediction[0] * 100:.1f}% Non-Fraud
+            </div>
+            <div style="width: {decrypted_prediction[1] * 100}%; background-color: red; color: white; text-align: center; padding: 5px 0;">
+                {decrypted_prediction[1] * 100:.1f}% Fraud
+            </div>
+        </div>
+        """
+        return "⚠️ Fraudulent ⚠️" if binary_prediction == 1 else "😊 Non-fraudulent 😊", gr.update(value="Decryption successful ✅"), bar_html
+        
+    except Exception as e:
+        logging.error(f"Error during prediction: {e}")
+        return "Error during prediction❌", "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.hex(), 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"