Update app.py

app.py

@@ -1,503 +1,763 @@
-"""A Gradio app for anonymizing text data using FHE."""
+# """A Gradio app for anonymizing text data using FHE."""
 
-import base64
-import os
-import re
-import subprocess
-import time
-import uuid
-from typing import Dict, List
+# import os
+# import re
+# import subprocess
+# import time
+# import uuid
+# from typing import Dict, List
 
-import gradio as gr
-import numpy
-import pandas as pd
-import requests
-from fhe_anonymizer import FHEAnonymizer
-from utils_demo import *
+# import numpy
+# import pandas as pd
+# import requests
+# from fhe_anonymizer import FHEAnonymizer
+# from utils_demo import *
 
-from concrete.ml.deployment import FHEModelClient
+# from concrete.ml.deployment import FHEModelClient
 
 
-# Ensure the directory is clean before starting processes or reading files
-clean_directory()  
 
-anonymizer = FHEAnonymizer()
+# import gradio as gr
+# from predictor import predict, key_already_generated, pre_process_encrypt_send_purchase, decrypt_prediction
+# import base64
 
-# Start the Uvicorn server hosting the FastAPI app
-subprocess.Popen(["uvicorn", "server:app"], cwd=CURRENT_DIR)
-time.sleep(3)
 
-# Load data from files required for the application
-UUID_MAP = read_json(MAPPING_UUID_PATH)
-ANONYMIZED_DOCUMENT = read_txt(ANONYMIZED_FILE_PATH)
-MAPPING_ANONYMIZED_SENTENCES = read_pickle(MAPPING_ANONYMIZED_SENTENCES_PATH)
-MAPPING_ENCRYPTED_SENTENCES = read_pickle(MAPPING_ENCRYPTED_SENTENCES_PATH)
-ORIGINAL_DOCUMENT = read_txt(ORIGINAL_FILE_PATH).split("\n\n")
-MAPPING_DOC_EMBEDDING = read_pickle(MAPPING_DOC_EMBEDDING_PATH)
+# # Ensure the directory is clean before starting processes or reading files
+# clean_directory()  
 
-print(f"{ORIGINAL_DOCUMENT=}\n")
-print(f"{MAPPING_DOC_EMBEDDING.keys()=}")
+# anonymizer = FHEAnonymizer()
 
-# 4. Data Processing and Operations (No specific operations shown here, assuming it's part of anonymizer or client usage)
+# # Start the Uvicorn server hosting the FastAPI app
+# subprocess.Popen(["uvicorn", "server:app"], cwd=CURRENT_DIR)
+# time.sleep(3)
 
-# 5. Utilizing External Services or APIs
-# (Assuming client initialization and anonymizer setup are parts of using external services or application-specific logic)
+# # Load data from files required for the application
+# UUID_MAP = read_json(MAPPING_UUID_PATH)
+# ANONYMIZED_DOCUMENT = read_txt(ANONYMIZED_FILE_PATH)
+# MAPPING_ANONYMIZED_SENTENCES = read_pickle(MAPPING_ANONYMIZED_SENTENCES_PATH)
+# MAPPING_ENCRYPTED_SENTENCES = read_pickle(MAPPING_ENCRYPTED_SENTENCES_PATH)
+# ORIGINAL_DOCUMENT = read_txt(ORIGINAL_FILE_PATH).split("\n\n")
+# MAPPING_DOC_EMBEDDING = read_pickle(MAPPING_DOC_EMBEDDING_PATH)
 
-# Generate a random user ID for this session
-USER_ID = numpy.random.randint(0, 2**32)
+# print(f"{ORIGINAL_DOCUMENT=}\n")
+# print(f"{MAPPING_DOC_EMBEDDING.keys()=}")
 
+# # 4. Data Processing and Operations (No specific operations shown here, assuming it's part of anonymizer or client usage)
 
-def select_static_anonymized_sentences_fn(selected_sentences: List):
+# # 5. Utilizing External Services or APIs
+# # (Assuming client initialization and anonymizer setup are parts of using external services or application-specific logic)
 
-    selected_sentences = [MAPPING_ANONYMIZED_SENTENCES[sentence] for sentence in selected_sentences]
+# # Generate a random user ID for this session
+# USER_ID = numpy.random.randint(0, 2**32)
 
-    anonymized_selected_sentence = sorted(selected_sentences, key=lambda x: x[0])
 
-    anonymized_selected_sentence = [sentence for _, sentence in anonymized_selected_sentence]
+# def select_static_anonymized_sentences_fn(selected_sentences: List):
 
-    return "\n\n".join(anonymized_selected_sentence)
+#     selected_sentences = [MAPPING_ANONYMIZED_SENTENCES[sentence] for sentence in selected_sentences]
 
+#     anonymized_selected_sentence = sorted(selected_sentences, key=lambda x: x[0])
 
-def key_gen_fn() -> Dict:
-    """Generate keys for a given user."""
+#     anonymized_selected_sentence = [sentence for _, sentence in anonymized_selected_sentence]
 
-    print("------------ Step 1: Key Generation:")
+#     return "\n\n".join(anonymized_selected_sentence)
 
-    print(f"Your user ID is: {USER_ID}....")
 
+# def key_gen_fn() -> Dict:
+#     """Generate keys for a given user."""
 
-    client = FHEModelClient(path_dir=DEPLOYMENT_DIR, key_dir=KEYS_DIR / f"{USER_ID}")
-    client.load()
+#     print("------------ Step 1: Key Generation:")
 
-    # Creates the private and evaluation keys on the client side
-    client.generate_private_and_evaluation_keys()
+#     print(f"Your user ID is: {USER_ID}....")
 
-    # Get the serialized evaluation keys
-    serialized_evaluation_keys = client.get_serialized_evaluation_keys()
-    assert isinstance(serialized_evaluation_keys, bytes)
 
-    # Save the evaluation key
-    evaluation_key_path = KEYS_DIR / f"{USER_ID}/evaluation_key"
+#     client = FHEModelClient(path_dir=DEPLOYMENT_DIR, key_dir=KEYS_DIR / f"{USER_ID}")
+#     client.load()
 
-    write_bytes(evaluation_key_path, serialized_evaluation_keys)
+#     # Creates the private and evaluation keys on the client side
+#     client.generate_private_and_evaluation_keys()
 
-    # anonymizer.generate_key()
+#     # Get the serialized evaluation keys
+#     serialized_evaluation_keys = client.get_serialized_evaluation_keys()
+#     assert isinstance(serialized_evaluation_keys, bytes)
 
-    if not evaluation_key_path.is_file():
-        error_message = (
-            f"Error Encountered While generating the evaluation {evaluation_key_path.is_file()=}"
-        )
-        print(error_message)
-        return {gen_key_btn: gr.update(value=error_message)}
-    else:
-        print("Keys have been generated ✅")
-        return {gen_key_btn: gr.update(value="Keys have been generated ✅")}
+#     # Save the evaluation key
+#     evaluation_key_path = KEYS_DIR / f"{USER_ID}/evaluation_key"
+
+#     write_bytes(evaluation_key_path, serialized_evaluation_keys)
 
+#     # anonymizer.generate_key()
 
-def encrypt_doc_fn(doc):
+#     if not evaluation_key_path.is_file():
+#         error_message = (
+#             f"Error Encountered While generating the evaluation {evaluation_key_path.is_file()=}"
+#         )
+#         print(error_message)
+#         return {gen_key_btn: gr.update(value=error_message)}
+#     else:
+#         print("Keys have been generated ✅")
+#         return {gen_key_btn: gr.update(value="Keys have been generated ✅")}
 
-    print(f"\n------------ Step 2.1: Doc encryption: {doc=}")
 
-    if not (KEYS_DIR / f"{USER_ID}/evaluation_key").is_file():
-        return {encrypted_doc_box: gr.update(value="Error ❌: Please generate the key first!", lines=10)}
+# def encrypt_doc_fn(doc):
 
-    # Retrieve the client API
-    client = FHEModelClient(path_dir=DEPLOYMENT_DIR, key_dir=KEYS_DIR / f"{USER_ID}")
-    client.load()
+#     print(f"\n------------ Step 2.1: Doc encryption: {doc=}")
 
-    encrypted_tokens = []
-    tokens = re.findall(r"(\b[\w\.\/\-@]+\b|[\s,.!?;:'\"-]+|\$\d+(?:\.\d+)?|\€\d+(?:\.\d+)?)", ' '.join(doc))
+#     if not (KEYS_DIR / f"{USER_ID}/evaluation_key").is_file():
+#         return {encrypted_doc_box: gr.update(value="Error ❌: Please generate the key first!", lines=10)}
+
+#     # Retrieve the client API
+#     client = FHEModelClient(path_dir=DEPLOYMENT_DIR, key_dir=KEYS_DIR / f"{USER_ID}")
+#     client.load()
+
+#     encrypted_tokens = []
+#     tokens = re.findall(r"(\b[\w\.\/\-@]+\b|[\s,.!?;:'\"-]+|\$\d+(?:\.\d+)?|\€\d+(?:\.\d+)?)", ' '.join(doc))
     
-    for token in tokens:
-        if token.strip() and re.match(r"\w+", token):
-            emb_x = MAPPING_DOC_EMBEDDING[token]
-            assert emb_x.shape == (1, 1024)
-            encrypted_x = client.quantize_encrypt_serialize(emb_x)
-            assert isinstance(encrypted_x, bytes)
-            encrypted_tokens.append(encrypted_x)
+#     for token in tokens:
+#         if token.strip() and re.match(r"\w+", token):
+#             emb_x = MAPPING_DOC_EMBEDDING[token]
+#             assert emb_x.shape == (1, 1024)
+#             encrypted_x = client.quantize_encrypt_serialize(emb_x)
+#             assert isinstance(encrypted_x, bytes)
+#             encrypted_tokens.append(encrypted_x)
 
-    print("Doc encrypted ✅ on Client Side")
+#     print("Doc encrypted ✅ on Client Side")
 
-    # No need to save it
-    # write_bytes(KEYS_DIR / f"{USER_ID}/encrypted_doc", b"".join(encrypted_tokens))
+#     # No need to save it
+#     # write_bytes(KEYS_DIR / f"{USER_ID}/encrypted_doc", b"".join(encrypted_tokens))
 
-    encrypted_quant_tokens_hex = [token.hex()[500:510] for token in encrypted_tokens]
+#     encrypted_quant_tokens_hex = [token.hex()[500:510] for token in encrypted_tokens]
 
-    return {
-        encrypted_doc_box: gr.update(value=" ".join(encrypted_quant_tokens_hex), lines=10),
-        anonymized_doc_output: gr.update(visible=True, value=None),
-    }
+#     return {
+#         encrypted_doc_box: gr.update(value=" ".join(encrypted_quant_tokens_hex), lines=10),
+#         anonymized_doc_output: gr.update(visible=True, value=None),
+#     }
     
 
-def encrypt_query_fn(query):
+# def encrypt_query_fn(query):
 
-    print(f"\n------------ Step 2: Query encryption: {query=}")
+#     print(f"\n------------ Step 2: Query encryption: {query=}")
 
-    if not (KEYS_DIR / f"{USER_ID}/evaluation_key").is_file():
-        return {output_encrypted_box: gr.update(value="Error ❌: Please generate the key first!", lines=8)}
+#     if not (KEYS_DIR / f"{USER_ID}/evaluation_key").is_file():
+#         return {output_encrypted_box: gr.update(value="Error ❌: Please generate the key first!", lines=8)}
 
-    if is_user_query_valid(query):
-        return {
-            query_box: gr.update(
-                value=(
-                    "Unable to process ❌: The request exceeds the length limit or falls "
-                    "outside the scope of this document. Please refine your query."
-                )
-            )
-        }
+#     if is_user_query_valid(query):
+#         return {
+#             query_box: gr.update(
+#                 value=(
+#                     "Unable to process ❌: The request exceeds the length limit or falls "
+#                     "outside the scope of this document. Please refine your query."
+#                 )
+#             )
+#         }
 
-    # Retrieve the client API
-    client = FHEModelClient(path_dir=DEPLOYMENT_DIR, key_dir=KEYS_DIR / f"{USER_ID}")
-    client.load()
+#     # Retrieve the client API
+#     client = FHEModelClient(path_dir=DEPLOYMENT_DIR, key_dir=KEYS_DIR / f"{USER_ID}")
+#     client.load()
 
-    encrypted_tokens = []
+#     encrypted_tokens = []
 
-    # Pattern to identify words and non-words (including punctuation, spaces, etc.)
-    tokens = re.findall(r"(\b[\w\.\/\-@]+\b|[\s,.!?;:'\"-]+)", query)
+#     # Pattern to identify words and non-words (including punctuation, spaces, etc.)
+#     tokens = re.findall(r"(\b[\w\.\/\-@]+\b|[\s,.!?;:'\"-]+)", query)
 
-    for token in tokens:
+#     for token in tokens:
 
-        # 1- Ignore non-words tokens
-        if bool(re.match(r"^\s+$", token)):
-            continue
+#         # 1- Ignore non-words tokens
+#         if bool(re.match(r"^\s+$", token)):
+#             continue
 
-        # 2- Directly append non-word tokens or whitespace to processed_tokens
+#         # 2- Directly append non-word tokens or whitespace to processed_tokens
 
-        # Prediction for each word
-        emb_x = get_batch_text_representation([token], EMBEDDINGS_MODEL, TOKENIZER)
-        encrypted_x = client.quantize_encrypt_serialize(emb_x)
-        assert isinstance(encrypted_x, bytes)
+#         # Prediction for each word
+#         emb_x = get_batch_text_representation([token], EMBEDDINGS_MODEL, TOKENIZER)
+#         encrypted_x = client.quantize_encrypt_serialize(emb_x)
+#         assert isinstance(encrypted_x, bytes)
 
-        encrypted_tokens.append(encrypted_x)
+#         encrypted_tokens.append(encrypted_x)
 
-    print("Data encrypted ✅ on Client Side")
+#     print("Data encrypted ✅ on Client Side")
 
-    assert len({len(token) for token in encrypted_tokens}) == 1
+#     assert len({len(token) for token in encrypted_tokens}) == 1
 
-    write_bytes(KEYS_DIR / f"{USER_ID}/encrypted_input", b"".join(encrypted_tokens))
-    write_bytes(
-        KEYS_DIR / f"{USER_ID}/encrypted_input_len", len(encrypted_tokens[0]).to_bytes(10, "big")
-    )
+#     write_bytes(KEYS_DIR / f"{USER_ID}/encrypted_input", b"".join(encrypted_tokens))
+#     write_bytes(
+#         KEYS_DIR / f"{USER_ID}/encrypted_input_len", len(encrypted_tokens[0]).to_bytes(10, "big")
+#     )
 
-    encrypted_quant_tokens_hex = [token.hex()[500:580] for token in encrypted_tokens]
+#     encrypted_quant_tokens_hex = [token.hex()[500:580] for token in encrypted_tokens]
 
-    return {
-        output_encrypted_box: gr.update(value=" ".join(encrypted_quant_tokens_hex), lines=8),
-        anonymized_query_output: gr.update(visible=True, value=None),
-        identified_words_output_df: gr.update(visible=False, value=None),
-    }
+#     return {
+#         output_encrypted_box: gr.update(value=" ".join(encrypted_quant_tokens_hex), lines=8),
+#         anonymized_query_output: gr.update(visible=True, value=None),
+#         identified_words_output_df: gr.update(visible=False, value=None),
+#     }
 
 
-def send_input_fn(query) -> Dict:
-    """Send the encrypted data and the evaluation key to the server."""
+# def send_input_fn(query) -> Dict:
+#     """Send the encrypted data and the evaluation key to the server."""
 
-    print("------------ Step 3.1: Send encrypted_data to the Server")
+#     print("------------ Step 3.1: Send encrypted_data to the Server")
 
-    evaluation_key_path = KEYS_DIR / f"{USER_ID}/evaluation_key"
-    encrypted_input_path = KEYS_DIR / f"{USER_ID}/encrypted_input"
-    encrypted_input_len_path = KEYS_DIR / f"{USER_ID}/encrypted_input_len"
+#     evaluation_key_path = KEYS_DIR / f"{USER_ID}/evaluation_key"
+#     encrypted_input_path = KEYS_DIR / f"{USER_ID}/encrypted_input"
+#     encrypted_input_len_path = KEYS_DIR / f"{USER_ID}/encrypted_input_len"
 
-    if not evaluation_key_path.is_file():
-        error_message = (
-            "Error Encountered While Sending Data to the Server: "
-            f"The key has been generated correctly - {evaluation_key_path.is_file()=}"
-        )
-        return {anonymized_query_output: gr.update(value=error_message)}
+#     if not evaluation_key_path.is_file():
+#         error_message = (
+#             "Error Encountered While Sending Data to the Server: "
+#             f"The key has been generated correctly - {evaluation_key_path.is_file()=}"
+#         )
+#         return {anonymized_query_output: gr.update(value=error_message)}
+
+#     if not encrypted_input_path.is_file():
+#         error_message = (
+#             "Error Encountered While Sending Data to the Server: The data has not been encrypted "
+#             f"correctly on the client side - {encrypted_input_path.is_file()=}"
+#         )
+#         return {anonymized_query_output: gr.update(value=error_message)}
+
+#     # Define the data and files to post
+#     data = {"user_id": USER_ID, "input": query}
+
+#     files = [
+#         ("files", open(evaluation_key_path, "rb")),
+#         ("files", open(encrypted_input_path, "rb")),
+#         ("files", open(encrypted_input_len_path, "rb")),
+#     ]
+
+#     # Send the encrypted input and evaluation key to the server
+#     url = SERVER_URL + "send_input"
+
+#     with requests.post(
+#         url=url,
+#         data=data,
+#         files=files,
+#     ) as resp:
+#         print("Data sent to the server ✅" if resp.ok else "Error ❌ in sending data to the server")
 
-    if not encrypted_input_path.is_file():
-        error_message = (
-            "Error Encountered While Sending Data to the Server: The data has not been encrypted "
-            f"correctly on the client side - {encrypted_input_path.is_file()=}"
-        )
-        return {anonymized_query_output: gr.update(value=error_message)}
 
-    # Define the data and files to post
-    data = {"user_id": USER_ID, "input": query}
+# def run_fhe_in_server_fn() -> Dict:
+#     """Run in FHE the anonymization of the query"""
+
+#     print("------------ Step 3.2: Run in FHE on the Server Side")
+
+#     evaluation_key_path = KEYS_DIR / f"{USER_ID}/evaluation_key"
+#     encrypted_input_path = KEYS_DIR / f"{USER_ID}/encrypted_input"
+
+#     if not evaluation_key_path.is_file():
+#         error_message = (
+#             "Error Encountered While Sending Data to the Server: "
+#             f"The key has been generated correctly - {evaluation_key_path.is_file()=}"
+#         )
+#         return {anonymized_query_output: gr.update(value=error_message)}
 
-    files = [
-        ("files", open(evaluation_key_path, "rb")),
-        ("files", open(encrypted_input_path, "rb")),
-        ("files", open(encrypted_input_len_path, "rb")),
-    ]
+#     if not encrypted_input_path.is_file():
+#         error_message = (
+#             "Error Encountered While Sending Data to the Server: The data has not been encrypted "
+#             f"correctly on the client side - {encrypted_input_path.is_file()=}"
+#         )
+#         return {anonymized_query_output: gr.update(value=error_message)}
+
+#     data = {
+#         "user_id": USER_ID,
+#     }
 
-    # Send the encrypted input and evaluation key to the server
-    url = SERVER_URL + "send_input"
+#     url = SERVER_URL + "run_fhe"
+
+#     with requests.post(
+#         url=url,
+#         data=data,
+#     ) as response:
+#         if not response.ok:
+#             return {
+#                 anonymized_query_output: gr.update(
+#                     value=(
+#                         "⚠️ An error occurred on the Server Side. "
+#                         "Please check connectivity and data transmission."
+#                     ),
+#                 ),
+#             }
+#         else:
+#             time.sleep(1)
+#             print(f"The query anonymization was computed in {response.json():.2f} s per token.")
+
+
+# def get_output_fn() -> Dict:
+
+#     print("------------ Step 3.3: Get the output from the Server Side")
+
+#     if not (KEYS_DIR / f"{USER_ID}/evaluation_key").is_file():
+#         error_message = (
+#             "Error Encountered While Sending Data to the Server: "
+#             "The key has not been generated correctly"
+#         )
+#         return {anonymized_query_output: gr.update(value=error_message)}
+
+#     if not (KEYS_DIR / f"{USER_ID}/encrypted_input").is_file():
+#         error_message = (
+#             "Error Encountered While Sending Data to the Server: "
+#             "The data has not been encrypted correctly on the client side"
+#         )
+#         return {anonymized_query_output: gr.update(value=error_message)}
+
+#     data = {
+#         "user_id": USER_ID,
+#     }
+
+#     # Retrieve the encrypted output
+#     url = SERVER_URL + "get_output"
+#     with requests.post(
+#         url=url,
+#         data=data,
+#     ) as response:
+#         if response.ok:
+#             print("Data received ✅ from the remote Server")
+#             response_data = response.json()
+#             encrypted_output_base64 = response_data["encrypted_output"]
+#             length_encrypted_output_base64 = response_data["length"]
 
-    with requests.post(
-        url=url,
-        data=data,
-        files=files,
-    ) as resp:
-        print("Data sent to the server ✅" if resp.ok else "Error ❌ in sending data to the server")
+#             # Decode the base64 encoded data
+#             encrypted_output = base64.b64decode(encrypted_output_base64)
+#             length_encrypted_output = base64.b64decode(length_encrypted_output_base64)
+
+#             # 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)
 
+#             write_bytes(CLIENT_DIR / f"{USER_ID}_encrypted_output", encrypted_output)
+#             write_bytes(CLIENT_DIR / f"{USER_ID}_encrypted_output_len", length_encrypted_output)
 
-def run_fhe_in_server_fn() -> Dict:
-    """Run in FHE the anonymization of the query"""
+#         else:
+#             print("Error ❌ in getting data to the server")
 
-    print("------------ Step 3.2: Run in FHE on the Server Side")
 
-    evaluation_key_path = KEYS_DIR / f"{USER_ID}/evaluation_key"
-    encrypted_input_path = KEYS_DIR / f"{USER_ID}/encrypted_input"
+# def decrypt_fn(text) -> Dict:
+#     """Dencrypt the data on the `Client Side`."""
 
-    if not evaluation_key_path.is_file():
-        error_message = (
-            "Error Encountered While Sending Data to the Server: "
-            f"The key has been generated correctly - {evaluation_key_path.is_file()=}"
-        )
-        return {anonymized_query_output: gr.update(value=error_message)}
+#     print("------------ Step 4: Dencrypt the data on the `Client Side`")
 
-    if not encrypted_input_path.is_file():
-        error_message = (
-            "Error Encountered While Sending Data to the Server: The data has not been encrypted "
-            f"correctly on the client side - {encrypted_input_path.is_file()=}"
-        )
-        return {anonymized_query_output: gr.update(value=error_message)}
+#     # Get the encrypted output path
+#     encrypted_output_path = CLIENT_DIR / f"{USER_ID}_encrypted_output"
 
-    data = {
-        "user_id": USER_ID,
-    }
+#     if not encrypted_output_path.is_file():
+#         error_message = """⚠️ Please ensure that: \n
+#                 - the connectivity \n
+#                 - the query has been submitted \n
+#                 - the evaluation key has been generated \n
+#                 - the server processed the encrypted data \n
+#                 - the Client received the data from the Server before decrypting the prediction
+#                 """
+#         print(error_message)
 
-    url = SERVER_URL + "run_fhe"
+#         return error_message, None
 
-    with requests.post(
-        url=url,
-        data=data,
-    ) as response:
-        if not response.ok:
-            return {
-                anonymized_query_output: gr.update(
-                    value=(
-                        "⚠️ An error occurred on the Server Side. "
-                        "Please check connectivity and data transmission."
-                    ),
-                ),
-            }
-        else:
-            time.sleep(1)
-            print(f"The query anonymization was computed in {response.json():.2f} s per token.")
+#     # Retrieve the client API
+#     client = FHEModelClient(path_dir=DEPLOYMENT_DIR, key_dir=KEYS_DIR / f"{USER_ID}")
+#     client.load()
+
+#     # Load the encrypted output as bytes
+#     encrypted_output = read_bytes(CLIENT_DIR / f"{USER_ID}_encrypted_output")
+#     length = int.from_bytes(read_bytes(CLIENT_DIR / f"{USER_ID}_encrypted_output_len"), "big")
 
+#     tokens = re.findall(r"(\b[\w\.\/\-@]+\b|[\s,.!?;:'\"-]+)", text)
+
+#     decrypted_output, identified_words_with_prob = [], []
 
-def get_output_fn() -> Dict:
+#     i = 0
+#     for token in tokens:
+
+#         # Directly append non-word tokens or whitespace to processed_tokens
+#         if bool(re.match(r"^\s+$", token)):
+#             continue
+#         else:
+#             encrypted_token = encrypted_output[i : i + length]
+#             prediction_proba = client.deserialize_decrypt_dequantize(encrypted_token)
+#             probability = prediction_proba[0][1]
+#             i += length
 
-    print("------------ Step 3.3: Get the output from the Server Side")
+#             if probability >= 0.77:
+#                 identified_words_with_prob.append((token, probability))
 
-    if not (KEYS_DIR / f"{USER_ID}/evaluation_key").is_file():
-        error_message = (
-            "Error Encountered While Sending Data to the Server: "
-            "The key has not been generated correctly"
-        )
-        return {anonymized_query_output: gr.update(value=error_message)}
+#                 # Use the existing UUID if available, otherwise generate a new one
+#                 tmp_uuid = UUID_MAP.get(token, str(uuid.uuid4())[:8])
+#                 decrypted_output.append(tmp_uuid)
+#                 UUID_MAP[token] = tmp_uuid
+#             else:
+#                 decrypted_output.append(token)
 
-    if not (KEYS_DIR / f"{USER_ID}/encrypted_input").is_file():
-        error_message = (
-            "Error Encountered While Sending Data to the Server: "
-            "The data has not been encrypted correctly on the client side"
-        )
-        return {anonymized_query_output: gr.update(value=error_message)}
-
-    data = {
-        "user_id": USER_ID,
-    }
-
-    # Retrieve the encrypted output
-    url = SERVER_URL + "get_output"
-    with requests.post(
-        url=url,
-        data=data,
-    ) as response:
-        if response.ok:
-            print("Data received ✅ from the remote Server")
-            response_data = response.json()
-            encrypted_output_base64 = response_data["encrypted_output"]
-            length_encrypted_output_base64 = response_data["length"]
-
-            # Decode the base64 encoded data
-            encrypted_output = base64.b64decode(encrypted_output_base64)
-            length_encrypted_output = base64.b64decode(length_encrypted_output_base64)
-
-            # 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)
-
-            write_bytes(CLIENT_DIR / f"{USER_ID}_encrypted_output", encrypted_output)
-            write_bytes(CLIENT_DIR / f"{USER_ID}_encrypted_output_len", length_encrypted_output)
-
-        else:
-            print("Error ❌ in getting data to the server")
-
-
-def decrypt_fn(text) -> Dict:
-    """Dencrypt the data on the `Client Side`."""
-
-    print("------------ Step 4: Dencrypt the data on the `Client Side`")
-
-    # Get the encrypted output path
-    encrypted_output_path = CLIENT_DIR / f"{USER_ID}_encrypted_output"
-
-    if not encrypted_output_path.is_file():
-        error_message = """⚠️ Please ensure that: \n
-                - the connectivity \n
-                - the query has been submitted \n
-                - the evaluation key has been generated \n
-                - the server processed the encrypted data \n
-                - the Client received the data from the Server before decrypting the prediction
-                """
-        print(error_message)
+#         # Update the UUID map with query.
+#         write_json(MAPPING_UUID_PATH, UUID_MAP)
 
-        return error_message, None
+#     # Removing Spaces Before Punctuation:
+#     anonymized_text = re.sub(r"\s([,.!?;:])", r"\1", " ".join(decrypted_output))
 
-    # Retrieve the client API
-    client = FHEModelClient(path_dir=DEPLOYMENT_DIR, key_dir=KEYS_DIR / f"{USER_ID}")
-    client.load()
+#     # Convert the list of identified words and probabilities into a DataFrame
+#     if identified_words_with_prob:
+#         identified_df = pd.DataFrame(
+#             identified_words_with_prob, columns=["Identified Words", "Probability"]
+#         )
+#     else:
+#         identified_df = pd.DataFrame(columns=["Identified Words", "Probability"])
 
-    # Load the encrypted output as bytes
-    encrypted_output = read_bytes(CLIENT_DIR / f"{USER_ID}_encrypted_output")
-    length = int.from_bytes(read_bytes(CLIENT_DIR / f"{USER_ID}_encrypted_output_len"), "big")
+#     print("Decryption done ✅ on Client Side")
 
-    tokens = re.findall(r"(\b[\w\.\/\-@]+\b|[\s,.!?;:'\"-]+)", text)
+#     return anonymized_text, identified_df
 
-    decrypted_output, identified_words_with_prob = [], []
 
-    i = 0
-    for token in tokens:
+# def anonymization_with_fn(selected_sentences, query):
 
-        # Directly append non-word tokens or whitespace to processed_tokens
-        if bool(re.match(r"^\s+$", token)):
-            continue
-        else:
-            encrypted_token = encrypted_output[i : i + length]
-            prediction_proba = client.deserialize_decrypt_dequantize(encrypted_token)
-            probability = prediction_proba[0][1]
-            i += length
+#     encrypt_query_fn(query)
 
-            if probability >= 0.77:
-                identified_words_with_prob.append((token, probability))
+#     send_input_fn(query)
 
-                # Use the existing UUID if available, otherwise generate a new one
-                tmp_uuid = UUID_MAP.get(token, str(uuid.uuid4())[:8])
-                decrypted_output.append(tmp_uuid)
-                UUID_MAP[token] = tmp_uuid
-            else:
-                decrypted_output.append(token)
+#     run_fhe_in_server_fn()
 
-        # Update the UUID map with query.
-        write_json(MAPPING_UUID_PATH, UUID_MAP)
+#     get_output_fn()
 
-    # Removing Spaces Before Punctuation:
-    anonymized_text = re.sub(r"\s([,.!?;:])", r"\1", " ".join(decrypted_output))
+#     anonymized_text, identified_df = decrypt_fn(query)
 
-    # Convert the list of identified words and probabilities into a DataFrame
-    if identified_words_with_prob:
-        identified_df = pd.DataFrame(
-            identified_words_with_prob, columns=["Identified Words", "Probability"]
-        )
-    else:
-        identified_df = pd.DataFrame(columns=["Identified Words", "Probability"])
+#     return {
+#         anonymized_doc_output: gr.update(value=select_static_anonymized_sentences_fn(selected_sentences)),
+#         anonymized_query_output: gr.update(value=anonymized_text),
+#         identified_words_output_df: gr.update(value=identified_df, visible=False),
+#     }
 
-    print("Decryption done ✅ on Client Side")
 
-    return anonymized_text, identified_df
+# def query_chatgpt_fn(anonymized_query, anonymized_document):
 
+#     print("------------ Step 5: ChatGPT communication")
 
-def anonymization_with_fn(selected_sentences, query):
+#     if not (KEYS_DIR / f"{USER_ID}/evaluation_key").is_file():
+#         error_message = "Error ❌: Please generate the key first!"
+#         return {chatgpt_response_anonymized: gr.update(value=error_message)}
 
-    encrypt_query_fn(query)
+#     if not (CLIENT_DIR / f"{USER_ID}_encrypted_output").is_file():
+#         error_message = "Error ❌: Please encrypt your query first!"
+#         return {chatgpt_response_anonymized: gr.update(value=error_message)}
 
-    send_input_fn(query)
+#     context_prompt = read_txt(PROMPT_PATH)
 
-    run_fhe_in_server_fn()
+#     # Prepare prompt
+#     query = (
+#         "Document content:\n```\n"
+#         + anonymized_document
+#         + "\n\n```"
+#         + "Query:\n```\n"
+#         + anonymized_query
+#         + "\n```"
+#     )
+#     print(f'Prompt of CHATGPT:\n{query}')
 
-    get_output_fn()
+#     completion = client.chat.completions.create(
+#         model="gpt-4-1106-preview",  # Replace with "gpt-4" if available
+#         messages=[
+#             {"role": "system", "content": context_prompt},
+#             {"role": "user", "content": query},
+#         ],
+#     )
+#     anonymized_response = completion.choices[0].message.content
+#     uuid_map = read_json(MAPPING_UUID_PATH)
 
-    anonymized_text, identified_df = decrypt_fn(query)
+#     inverse_uuid_map = {
+#         v: k for k, v in uuid_map.items()
+#     }  # TODO load the inverse mapping from disk for efficiency
 
-    return {
-        anonymized_doc_output: gr.update(value=select_static_anonymized_sentences_fn(selected_sentences)),
-        anonymized_query_output: gr.update(value=anonymized_text),
-        identified_words_output_df: gr.update(value=identified_df, visible=False),
-    }
+#     # Pattern to identify words and non-words (including punctuation, spaces, etc.)
+#     tokens = re.findall(r"(\b[\w\.\/\-@]+\b|[\s,.!?;:'\"-]+)", anonymized_response)
+#     processed_tokens = []
 
+#     for token in tokens:
+#         # Directly append non-word tokens or whitespace to processed_tokens
+#         if not token.strip() or not re.match(r"\w+", token):
+#             processed_tokens.append(token)
+#             continue
 
-def query_chatgpt_fn(anonymized_query, anonymized_document):
+#         if token in inverse_uuid_map:
+#             processed_tokens.append(inverse_uuid_map[token])
+#         else:
+#             processed_tokens.append(token)
+#     deanonymized_response = "".join(processed_tokens)
 
-    print("------------ Step 5: ChatGPT communication")
+#     return {chatgpt_response_anonymized: gr.update(value=anonymized_response), 
+#             chatgpt_response_deanonymized: gr.update(value=deanonymized_response)}
 
-    if not (KEYS_DIR / f"{USER_ID}/evaluation_key").is_file():
-        error_message = "Error ❌: Please generate the key first!"
-        return {chatgpt_response_anonymized: gr.update(value=error_message)}
 
-    if not (CLIENT_DIR / f"{USER_ID}_encrypted_output").is_file():
-        error_message = "Error ❌: Please encrypt your query first!"
-        return {chatgpt_response_anonymized: gr.update(value=error_message)}
+# demo = gr.Blocks(css=".markdown-body { font-size: 18px; }")
 
-    context_prompt = read_txt(PROMPT_PATH)
+# with demo:
 
-    # Prepare prompt
-    query = (
-        "Document content:\n```\n"
-        + anonymized_document
-        + "\n\n```"
-        + "Query:\n```\n"
-        + anonymized_query
-        + "\n```"
-    )
-    print(f'Prompt of CHATGPT:\n{query}')
-
-    completion = client.chat.completions.create(
-        model="gpt-4-1106-preview",  # Replace with "gpt-4" if available
-        messages=[
-            {"role": "system", "content": context_prompt},
-            {"role": "user", "content": query},
-        ],
-    )
-    anonymized_response = completion.choices[0].message.content
-    uuid_map = read_json(MAPPING_UUID_PATH)
+#     gr.Markdown(
+#         """
+#         <p align="center">
+#             <img width=200 src="https://user-images.githubusercontent.com/5758427/197816413-d9cddad3-ba38-4793-847d-120975e1da11.png">
+#         </p>
+#         """)
+        
+#     gr.Markdown(
+#         """
+#         <h1 style="text-align: center;">Encrypted Anonymization Using Fully Homomorphic Encryption</h1>
+#         <p align="center">
+#             <a href="https://github.com/zama-ai/concrete-ml"> <img style="vertical-align: middle; display:inline-block; margin-right: 3px;" width=15 src="file/images/logos/github.png">Concrete-ML</a>
+#             —
+#             <a href="https://docs.zama.ai/concrete-ml"> <img style="vertical-align: middle; display:inline-block; margin-right: 3px;" width=15 src="file/images/logos/documentation.png">Documentation</a>
+#             —
+#             <a href=" https://community.zama.ai/c/concrete-ml/8"> <img style="vertical-align: middle; display:inline-block; margin-right: 3px;" width=15 src="file/images/logos/community.png">Community</a>
+#             —
+#             <a href="https://twitter.com/zama_fhe"> <img style="vertical-align: middle; display:inline-block; margin-right: 3px;" width=15 src="file/images/logos/x.png">@zama_fhe</a>
+#         </p>
+#         """
+#     )
+
+#     gr.Markdown(
+#     """
+#     <p align="center" style="font-size: 16px;">
+#         Anonymization is the process of removing personally identifiable information (PII) data from 
+#         a document in order to protect individual privacy.</p>
+
+#     <p align="center" style="font-size: 16px;">
+#         Encrypted anonymization uses Fully Homomorphic Encryption (FHE) to anonymize personally 
+#         identifiable information (PII) within encrypted documents, enabling computations to be 
+#         performed on the encrypted data.</p>
+
+#     <p align="center" style="font-size: 16px;">
+#         In the example above, we're showing how encrypted anonymization can be leveraged to use LLM 
+#         services such as ChatGPT in a privacy-preserving manner.</p>
+#     """
+#     )    
+    
+#     # gr.Markdown(
+#     #     """
+#     #     <p align="center">
+#     #         <img width="75%" height="30%" src="https://raw.githubusercontent.com/kcelia/Img/main/fhe_anonymization_banner.png">
+#     #     </p>
+#     #     """
+#     # )
+#     gr.Markdown(
+#         f"""
+#         <p align="center">
+#             <img width="75%" height="30%" src="https://huggingface.co/spaces/Tenefix/private-fhe-fraud-detection/resolve/main/Img/schema.png">
+#         </p>
+#         """
+#     )
+
+
+#     ########################## Key Gen Part ##########################
+
+#     gr.Markdown(
+#         "## Step 1: Generate the keys\n\n"
+#         """In Fully Homomorphic Encryption (FHE) methods, two types of keys are created. The first 
+#         type, called secret keys, are used to encrypt and decrypt the user's data. The second type, 
+#         called evaluation keys, enables a server to work on the encrypted data without seeing the 
+#         actual data.
+#         """
+#     )
+
+#     gen_key_btn = gr.Button("Generate the secret and evaluation keys")
+
+#     gen_key_btn.click(
+#         key_gen_fn,
+#         inputs=[],
+#         outputs=[gen_key_btn],
+#     )
+
+#     ########################## Main document Part ##########################
+
+#     gr.Markdown("<hr />")
+#     gr.Markdown("## Step 2.1: Select the document you want to encrypt\n\n"
+#         """To make it simple, we pre-compiled the following document, but you are free to choose 
+#         on which part you want to run this example.
+#         """
+#     )
+
+#     with gr.Row():
+#         with gr.Column(scale=5):
+#             original_sentences_box = gr.CheckboxGroup(
+#                 ORIGINAL_DOCUMENT,
+#                 value=ORIGINAL_DOCUMENT,
+#                 label="Contract:",
+#                 show_label=True,
+#             )
+
+#         with gr.Column(scale=1, min_width=6):
+#             gr.HTML("<div style='height: 77px;'></div>")
+#             encrypt_doc_btn = gr.Button("Encrypt the document")
+
+#         with gr.Column(scale=5):
+#             encrypted_doc_box = gr.Textbox(
+#                 label="Encrypted document:", show_label=True, interactive=False, lines=10
+#             )
+
+
+#     ########################## User Query Part ##########################
+
+#     gr.Markdown("<hr />")
+#     gr.Markdown("## Step 2.2: Select the prompt you want to encrypt\n\n"
+#         """Please choose from the predefined options in 
+#         <span style='color:grey'>“Prompt examples”</span> or craft a custom question in 
+#         the <span style='color:grey'>“Customized prompt”</span> text box.        
+#         Remain concise and relevant to the context. Any off-topic query will not be processed.""")
+
+#     with gr.Row():
+#         with gr.Column(scale=5):
+
+#             with gr.Column(scale=5):
+#                 default_query_box = gr.Dropdown(
+#                     list(DEFAULT_QUERIES.values()), label="PROMPT EXAMPLES:"
+#                 )
+
+#             gr.Markdown("Or")
+
+#             query_box = gr.Textbox(
+#                 value="What is Kate international bank account number?", label="CUSTOMIZED PROMPT:", interactive=True
+#             )
+
+#             default_query_box.change(
+#                 fn=lambda default_query_box: default_query_box,
+#                 inputs=[default_query_box],
+#                 outputs=[query_box],
+#             )
+
+#         with gr.Column(scale=1, min_width=6):
+#             gr.HTML("<div style='height: 77px;'></div>")
+#             encrypt_query_btn = gr.Button("Encrypt the prompt")
+#             # gr.HTML("<div style='height: 50px;'></div>")
+
+#         with gr.Column(scale=5):
+#             output_encrypted_box = gr.Textbox(
+#                 label="Encrypted anonymized query that will be sent to the anonymization server:",
+#                 lines=8,
+#             )
+
+#     ########################## FHE processing Part ##########################
+
+#     gr.Markdown("<hr />")
+#     gr.Markdown("## Step 3: Anonymize the document and the prompt using FHE")
+#     gr.Markdown(
+#         """Once the client encrypts the document and the prompt locally, it will be sent to a remote 
+#         server to perform the anonymization on encrypted data. When the computation is done, the 
+#         server will return the result to the client for decryption.
+#         """
+#     )
+
+#     run_fhe_btn = gr.Button("Anonymize using FHE")
+
+#     with gr.Row():
+#         with gr.Column(scale=5):
+
+#             anonymized_doc_output = gr.Textbox(
+#                 label="Decrypted and anonymized document", lines=10, interactive=True
+#             )
+
+#         with gr.Column(scale=5):
+
+#             anonymized_query_output = gr.Textbox(
+#                 label="Decrypted and anonymized prompt", lines=10, interactive=True
+#             )
+
+
+#     identified_words_output_df = gr.Dataframe(label="Identified words:", visible=False)
+
+#     encrypt_doc_btn.click(
+#         fn=encrypt_doc_fn,
+#         inputs=[original_sentences_box],
+#         outputs=[encrypted_doc_box, anonymized_doc_output],
+#     )
 
-    inverse_uuid_map = {
-        v: k for k, v in uuid_map.items()
-    }  # TODO load the inverse mapping from disk for efficiency
+#     encrypt_query_btn.click(
+#         fn=encrypt_query_fn,
+#         inputs=[query_box],
+#         outputs=[
+#             query_box,
+#             output_encrypted_box,
+#             anonymized_query_output,
+#             identified_words_output_df,
+#         ],
+#     )
 
-    # Pattern to identify words and non-words (including punctuation, spaces, etc.)
-    tokens = re.findall(r"(\b[\w\.\/\-@]+\b|[\s,.!?;:'\"-]+)", anonymized_response)
-    processed_tokens = []
+#     run_fhe_btn.click(
+#         anonymization_with_fn,
+#         inputs=[original_sentences_box, query_box],
+#         outputs=[anonymized_doc_output, anonymized_query_output, identified_words_output_df],
+#     )
 
-    for token in tokens:
-        # Directly append non-word tokens or whitespace to processed_tokens
-        if not token.strip() or not re.match(r"\w+", token):
-            processed_tokens.append(token)
-            continue
+#     ########################## ChatGpt Part ##########################
 
-        if token in inverse_uuid_map:
-            processed_tokens.append(inverse_uuid_map[token])
-        else:
-            processed_tokens.append(token)
-    deanonymized_response = "".join(processed_tokens)
+#     gr.Markdown("<hr />")
+#     gr.Markdown("## Step 4: Send anonymized prompt to ChatGPT")
+#     gr.Markdown(
+#         """After securely anonymizing the query with FHE, 
+#         you can forward it to ChatGPT without having any concern about information leakage."""
+#     )
 
-    return {chatgpt_response_anonymized: gr.update(value=anonymized_response), 
-            chatgpt_response_deanonymized: gr.update(value=deanonymized_response)}
+#     chatgpt_button = gr.Button("Query ChatGPT")
 
+#     with gr.Row():
+#         chatgpt_response_anonymized = gr.Textbox(label="ChatGPT's anonymized response:", lines=5)
+#         chatgpt_response_deanonymized = gr.Textbox(
+#             label="ChatGPT's non-anonymized response:", lines=5
+#         )
 
+#     chatgpt_button.click(
+#         query_chatgpt_fn,
+#         inputs=[anonymized_query_output, anonymized_doc_output],
+#         outputs=[chatgpt_response_anonymized, chatgpt_response_deanonymized],
+#     )
+
+#     gr.Markdown(
+#         """**Please note**: As this space is intended solely for demonstration purposes, some 
+#         private information may be missed during by the anonymization algorithm. Please validate the 
+#         following query before sending it to ChatGPT."""
+#     )
+# # Launch the app
+# demo.launch(share=False)
+
+
+
+
+
+
+
+import gradio as gr
+from predictor import predict, key_already_generated, pre_process_encrypt_send_purchase, decrypt_prediction
+import base64
+
+def key_generated():
+    """
+    Check if the evaluation keys have already been generated.
+    Returns:
+        bool: True if the evaluation keys have already been generated, False otherwise.
+    """
+    if not key_already_generated():
+        error_message = (
+            f"Error Encountered While generating the evaluation keys."
+        )
+        print(error_message)
+        return {gen_key_btn: gr.update(value=error_message)}
+    else:
+        print("Keys have been generated ✅")
+        return {gen_key_btn: gr.update(value="Keys have been generated ✅")}
+
+    
 demo = gr.Blocks(css=".markdown-body { font-size: 18px; }")
 
 with demo:
+    with gr.Row():
+        with gr.Column(elem_id="center_column"):
+            gr.Image("Img/zama.png", width=200, show_label=False)
+        with gr.Column(elem_id="center_column"):
+            gr.Image("Img/Epita.png", width=200, show_label=False)
+
+
 
     gr.Markdown(
         """
+        <h1 style="text-align: center;">Fraud Detection with FHE Model</h1>
         <p align="center">
-            <img width=200 src="https://user-images.githubusercontent.com/5758427/197816413-d9cddad3-ba38-4793-847d-120975e1da11.png">
-        </p>
-        """)
-        
-    gr.Markdown(
-        """
-        <h1 style="text-align: center;">Encrypted Anonymization Using Fully Homomorphic Encryption</h1>
-        <p align="center">
-            <a href="https://github.com/zama-ai/concrete-ml"> <img style="vertical-align: middle; display:inline-block; margin-right: 3px;" width=15 src="file/images/logos/github.png">Concrete-ML</a>
+            <a href="https://github.com/CirSandro/private-fhe-fraud-detection">
+                <span style="vertical-align: middle; display:inline-block; margin-right: 3px;">💳</span>private-fhe-fraud-detection
+            </a>
             —
-            <a href="https://docs.zama.ai/concrete-ml"> <img style="vertical-align: middle; display:inline-block; margin-right: 3px;" width=15 src="file/images/logos/documentation.png">Documentation</a>
-            —
-            <a href=" https://community.zama.ai/c/concrete-ml/8"> <img style="vertical-align: middle; display:inline-block; margin-right: 3px;" width=15 src="file/images/logos/community.png">Community</a>
-            —
-            <a href="https://twitter.com/zama_fhe"> <img style="vertical-align: middle; display:inline-block; margin-right: 3px;" width=15 src="file/images/logos/x.png">@zama_fhe</a>
+            <a href="https://docs.zama.ai/concrete-ml">
+                <span style="vertical-align: middle; display:inline-block; margin-right: 3px;">🔒</span>Documentation Concrete-ML
+            </a>
         </p>
         """
     )
@@ -505,35 +765,61 @@ with demo:
     gr.Markdown(
     """
     <p align="center" style="font-size: 16px;">
-        Anonymization is the process of removing personally identifiable information (PII) data from 
-        a document in order to protect individual privacy.</p>
+        How to detect bank fraud without using your personal data ?</p>
+    """
+    )
 
-    <p align="center" style="font-size: 16px;">
-        Encrypted anonymization uses Fully Homomorphic Encryption (FHE) to anonymize personally 
-        identifiable information (PII) within encrypted documents, enabling computations to be 
-        performed on the encrypted data.</p>
+    with gr.Accordion("What is bank fraud detection?", open=False):
+        gr.Markdown(
+            """
+            Bank fraud detection is the process of identifying fraudulent activities or transactions 
+            that may pose a risk to a bank or its customers. It is essential to detect fraudulent 
+            activities to prevent financial losses and protect the integrity of the banking system.
+            """
+        )
+
+    with gr.Accordion("Why is it important to protect this data?", open=False):
+        gr.Markdown(
+            """
+            Banking and financial data often contain sensitive personal information, such as income,
+            spending habits, and account numbers. Protecting this information ensures that customers'
+            privacy is respected and safeguarded from unauthorized access.
+            """
+        )
+
+    with gr.Accordion("Why is Fully Homomorphic Encryption (FHE) a good solution?", open=False):
+        gr.Markdown(
+            """
+            Fully Homomorphic Encryption (FHE) is a powerful technique for enhancing privacy and accuracy
+            in the context of fraud detection, particularly when dealing with sensitive banking data. FHE
+            allows for the encryption of data, which can then be processed and analyzed without ever needing
+            to decrypt it. 
+            Each party involved in the detection process can collaborate without compromising user privacy,
+            minimizing the risk of data leaks or breaches. The data remains confidential throughout the entire
+            process, ensuring that the privacy of users is maintained.
+            """
+        )
+
+    gr.Markdown(
+        """
+        <p style="text-align: center;">
+            Below, we will explain the flow in the image by simulating a purchase you've just made, and show you how our fraud detection model processes the transaction.
+        </p>
+        """
+    )
 
-    <p align="center" style="font-size: 16px;">
-        In the example above, we're showing how encrypted anonymization can be leveraged to use LLM 
-        services such as ChatGPT in a privacy-preserving manner.</p>
-    """
-    )    
-    
     # gr.Markdown(
-    #     """
+    #     f"""
     #     <p align="center">
-    #         <img width="75%" height="30%" src="https://raw.githubusercontent.com/kcelia/Img/main/fhe_anonymization_banner.png">
+    #         <img width="75%" height="30%" src="https://huggingface.co/spaces/Tenefix/private-fhe-fraud-detection/resolve/main/Img/schema.png">
     #     </p>
     #     """
     # )
-    gr.Markdown(
-        f"""
-        <p align="center">
-            <img width="75%" height="30%" src="https://huggingface.co/spaces/Tenefix/private-fhe-fraud-detection/resolve/main/Img/schema.png">
-        </p>
-        """
-    )
+    with gr.Row():
+        with gr.Column(elem_id="center_column"):
+            gr.Image("Img/schema.png", width=200, show_label=False)
 
+    gr.Markdown("<hr />")
 
     ########################## Key Gen Part ##########################
 
@@ -549,158 +835,178 @@ with demo:
     gen_key_btn = gr.Button("Generate the secret and evaluation keys")
 
     gen_key_btn.click(
-        key_gen_fn,
+        key_generated,
         inputs=[],
         outputs=[gen_key_btn],
-    )
-
-    ########################## Main document Part ##########################
+    )#547
 
     gr.Markdown("<hr />")
-    gr.Markdown("## Step 2.1: Select the document you want to encrypt\n\n"
-        """To make it simple, we pre-compiled the following document, but you are free to choose 
-        on which part you want to run this example.
-        """
-    )
-
-    with gr.Row():
-        with gr.Column(scale=5):
-            original_sentences_box = gr.CheckboxGroup(
-                ORIGINAL_DOCUMENT,
-                value=ORIGINAL_DOCUMENT,
-                label="Contract:",
-                show_label=True,
-            )
 
-        with gr.Column(scale=1, min_width=6):
-            gr.HTML("<div style='height: 77px;'></div>")
-            encrypt_doc_btn = gr.Button("Encrypt the document")
+    ########################## Encrypt Data ##########################
 
-        with gr.Column(scale=5):
-            encrypted_doc_box = gr.Textbox(
-                label="Encrypted document:", show_label=True, interactive=False, lines=10
-            )
-
-
-    ########################## User Query Part ##########################
+    gr.Markdown(
+        "## Step 2: Make your purchase\n\n"
+        """ 
+        🛍️ It's time to shop! To simulate your latest purchase, please provide the details of your most recent transaction.
+        
+        If you don't have an idea, you can pre-fill with an example of fraud or non-fraud.
+        """
+    )
 
-    gr.Markdown("<hr />")
-    gr.Markdown("## Step 2.2: Select the prompt you want to encrypt\n\n"
-        """Please choose from the predefined options in 
-        <span style='color:grey'>“Prompt examples”</span> or craft a custom question in 
-        the <span style='color:grey'>“Customized prompt”</span> text box.        
-        Remain concise and relevant to the context. Any off-topic query will not be processed.""")
+    def prefill_fraud():
+        return 34, 50, 3, False, False, False, True
+    
+    def prefill_no_fraud():
+        return 12, 2, 0.7, True, False, True, False
 
     with gr.Row():
-        with gr.Column(scale=5):
-
-            with gr.Column(scale=5):
-                default_query_box = gr.Dropdown(
-                    list(DEFAULT_QUERIES.values()), label="PROMPT EXAMPLES:"
-                )
+        prefill_button = gr.Button("Exemple Fraud")
+        prefill_button_no = gr.Button("Exemple No-Fraud")
 
-            gr.Markdown("Or")
-
-            query_box = gr.Textbox(
-                value="What is Kate international bank account number?", label="CUSTOMIZED PROMPT:", interactive=True
+    with gr.Row():
+        with gr.Column():
+            distance_home = gr.Number(
+                minimum=float(0),
+                maximum=float(22000),
+                step=1, 
+                value=10,
+                label="Distance from Home", 
+                info="How far was the purchase from your home (in km)?"
             )
-
-            default_query_box.change(
-                fn=lambda default_query_box: default_query_box,
-                inputs=[default_query_box],
-                outputs=[query_box],
+            distance_last = gr.Number(
+                minimum=float(0),
+                maximum=float(22000),
+                step=1,
+                value=1,
+                label="Distance from Last Transaction", 
+                info="Distance between this purchase and the last one (in km)?"
             )
-
-        with gr.Column(scale=1, min_width=6):
-            gr.HTML("<div style='height: 77px;'></div>")
-            encrypt_query_btn = gr.Button("Encrypt the prompt")
-            # gr.HTML("<div style='height: 50px;'></div>")
-
-        with gr.Column(scale=5):
-            output_encrypted_box = gr.Textbox(
-                label="Encrypted anonymized query that will be sent to the anonymization server:",
-                lines=8,
+            ratio = gr.Number(
+                minimum=float(0),
+                maximum=float(10000),
+                step=0.1,
+                value=1,
+                label="Ratio to Median Purchase Price", 
+                info="Purchase ratio compared to your average purchase",
+            )
+            repeat_retailer = gr.Checkbox(
+                label="Repeat Retailer",
+                info="Check if you are purchasing from the same retailer as your last transaction"
+            )
+            used_chip = gr.Checkbox(
+                label="Used Chip",
+                info="Check if you used a chip card for this transaction"
+            )
+            used_pin_number = gr.Checkbox(
+                label="Used Pin Number",
+                info="Check if you used your PIN number during the transaction"
+            )
+            online = gr.Checkbox(
+                label="Online Order",
+                info="Check if you made your purchase online"
             )
 
-    ########################## FHE processing Part ##########################
-
-    gr.Markdown("<hr />")
-    gr.Markdown("## Step 3: Anonymize the document and the prompt using FHE")
-    gr.Markdown(
-        """Once the client encrypts the document and the prompt locally, it will be sent to a remote 
-        server to perform the anonymization on encrypted data. When the computation is done, the 
-        server will return the result to the client for decryption.
-        """
+        
+    prefill_button.click(
+        fn=prefill_fraud,
+        inputs=[],
+        outputs=[
+            distance_home, 
+            distance_last, 
+            ratio, 
+            repeat_retailer, 
+            used_chip, 
+            used_pin_number, 
+            online
+        ]
     )
 
-    run_fhe_btn = gr.Button("Anonymize using FHE")
+    prefill_button_no.click(
+        fn=prefill_no_fraud,
+        inputs=[],
+        outputs=[
+            distance_home, 
+            distance_last, 
+            ratio, 
+            repeat_retailer, 
+            used_chip, 
+            used_pin_number, 
+            online
+        ]
+    )
 
     with gr.Row():
-        with gr.Column(scale=5):
-
-            anonymized_doc_output = gr.Textbox(
-                label="Decrypted and anonymized document", lines=10, interactive=True
-            )
-
-        with gr.Column(scale=5):
-
-            anonymized_query_output = gr.Textbox(
-                label="Decrypted and anonymized prompt", lines=10, interactive=True
+        with gr.Column(scale=2):
+            encrypt_button_applicant = gr.Button("Encrypt the inputs and send to server.")
+            
+            encrypted_input_applicant = gr.Textbox(
+                label="Encrypted input representation:", max_lines=2, interactive=False
             )
 
+    encrypt_button_applicant.click(
+        pre_process_encrypt_send_purchase,
+        inputs=[distance_home, distance_last, ratio, repeat_retailer, used_chip, used_pin_number, \
+                online],
+        outputs=[encrypted_input_applicant, encrypt_button_applicant],
+    )
 
-    identified_words_output_df = gr.Dataframe(label="Identified words:", visible=False)
+    gr.Markdown("<hr />")
 
-    encrypt_doc_btn.click(
-        fn=encrypt_doc_fn,
-        inputs=[original_sentences_box],
-        outputs=[encrypted_doc_box, anonymized_doc_output],
-    )
+    ########################## Model Prediction ##########################
 
-    encrypt_query_btn.click(
-        fn=encrypt_query_fn,
-        inputs=[query_box],
-        outputs=[
-            query_box,
-            output_encrypted_box,
-            anonymized_query_output,
-            identified_words_output_df,
-        ],
+    gr.Markdown("## Step 3: Run the FHE evaluation.")
+    gr.Markdown("<span style='color:grey'>Server Side</span>")
+    gr.Markdown(
+        """
+        It's high time to launch our prediction, by pressing the button you will launch the 
+        fraud analysis that our fictitious bank offers you.
+        This server employs a [Random Forest (by Concrete-ML)](https://github.com/zama-ai/concrete-ml/blob/release/1.8.x/docs/references/api/concrete.ml.sklearn.rf.md#class-randomforestclassifier)
+        classifier model that has been trained on a synthetic data-set.
+        """
     )
 
-    run_fhe_btn.click(
-        anonymization_with_fn,
-        inputs=[original_sentences_box, query_box],
-        outputs=[anonymized_doc_output, anonymized_query_output, identified_words_output_df],
+    execute_fhe_button = gr.Button("Run the FHE evaluation.")
+    fhe_execution_time = gr.Textbox(
+        label="Total FHE execution time (in seconds):", max_lines=1, interactive=False
     )
 
-    ########################## ChatGpt Part ##########################
+    # Button to send the encodings to the server using post method
+    execute_fhe_button.click(predict, outputs=[fhe_execution_time, execute_fhe_button])
 
     gr.Markdown("<hr />")
-    gr.Markdown("## Step 4: Send anonymized prompt to ChatGPT")
+
+    ########################## Decrypt Prediction ##########################
+
+    gr.Markdown("## Step 4: Receive the encrypted output from the server and decrypt.")
     gr.Markdown(
-        """After securely anonymizing the query with FHE, 
-        you can forward it to ChatGPT without having any concern about information leakage."""
+        """
+        🔔 You will receive a notification! Is this a Fraud? The message is decrypted by pressing the button.
+        """
     )
 
-    chatgpt_button = gr.Button("Query ChatGPT")
+    get_output_button = gr.Button("Decrypt the prediction.")
+    prediction_output = gr.Textbox(
+        label="Prediction", max_lines=1, interactive=False
+    )
+    prediction_bar = gr.HTML(label="Prediction Bar")  # For the percentage bar
 
-    with gr.Row():
-        chatgpt_response_anonymized = gr.Textbox(label="ChatGPT's anonymized response:", lines=5)
-        chatgpt_response_deanonymized = gr.Textbox(
-            label="ChatGPT's non-anonymized response:", lines=5
-        )
+    get_output_button.click(
+        decrypt_prediction,
+        outputs=[prediction_output, get_output_button, prediction_bar],
+    )
+    
 
-    chatgpt_button.click(
-        query_chatgpt_fn,
-        inputs=[anonymized_query_output, anonymized_doc_output],
-        outputs=[chatgpt_response_anonymized, chatgpt_response_deanonymized],
+    gr.Markdown(
+        """
+        You now know that it is possible to detect bank fraud without knowing your personal information.
+        """
     )
 
     gr.Markdown(
-        """**Please note**: As this space is intended solely for demonstration purposes, some 
-        private information may be missed during by the anonymization algorithm. Please validate the 
-        following query before sending it to ChatGPT."""
+        "The app was built with [Concrete-ML](https://github.com/zama-ai/concrete-ml), a "
+        "Privacy-Preserving Machine Learning (PPML) open-source set of tools by [Zama](https://zama.ai/). "
+        "Try it yourself and don't forget to star on Github &#11088;."
     )
-# Launch the app
-demo.launch(share=False)
+
+if __name__ == "__main__":
+    demo.launch()