maheshdev209/fhehp
0
1import subprocess2import time3from typing import Dict, List, Tuple4 5import gradio as gr # pylint: disable=import-error6import numpy as np7import pandas as pd8import requests9from symptoms_categories import SYMPTOMS_LIST10from utils import (11 CLIENT_DIR,12 CURRENT_DIR,13 DEPLOYMENT_DIR,14 INPUT_BROWSER_LIMIT,15 KEYS_DIR,16 SERVER_URL,17 TARGET_COLUMNS,18 TRAINING_FILENAME,19 clean_directory,20 get_disease_name,21 load_data,22 pretty_print,23)24 25from concrete.ml.deployment import FHEModelClient26 27subprocess.Popen(["uvicorn", "server:app"], cwd=CURRENT_DIR)28time.sleep(3)29 30# pylint: disable=c-extension-no-member,invalid-name31 32 33def is_none(obj) -> bool:34 """35 Check if the object is None.36 37 Args:38 obj (any): The input to be checked.39 40 Returns:41 bool: True if the object is None or empty, False otherwise.42 """43 return obj is None or (obj is not None and len(obj) < 1)44 45 46def display_default_symptoms_fn(default_disease: str) -> Dict:47 """48 Displays the symptoms of a given existing disease.49 50 Args:51 default_disease (str): Disease52 Returns:53 Dict: The according symptoms54 """55 df = pd.read_csv(TRAINING_FILENAME)56 df_filtred = df[df[TARGET_COLUMNS[1]] == default_disease]57 58 return {59 default_symptoms: gr.update(60 visible=True,61 value=pretty_print(62 df_filtred.columns[df_filtred.eq(1).any()].to_list(), delimiter=", "63 ),64 )65 }66 67 68def get_user_symptoms_from_checkboxgroup(checkbox_symptoms: List) -> np.array:69 """70 Convert the user symptoms into a binary vector representation.71 72 Args:73 checkbox_symptoms (List): A list of user symptoms.74 75 Returns:76 np.array: A binary vector representing the user's symptoms.77 78 Raises:79 KeyError: If a provided symptom is not recognized as a valid symptom.80 81 """82 symptoms_vector = {key: 0 for key in valid_symptoms}83 for pretty_symptom in checkbox_symptoms:84 original_symptom = "_".join((pretty_symptom.lower().split(" ")))85 if original_symptom not in symptoms_vector.keys():86 raise KeyError(87 f"The symptom '{original_symptom}' you provided is not recognized as a valid "88 f"symptom.\nHere is the list of valid symptoms: {symptoms_vector}"89 )90 symptoms_vector[original_symptom] = 191 92 user_symptoms_vect = np.fromiter(symptoms_vector.values(), dtype=float)[np.newaxis, :]93 94 assert all(value == 0 or value == 1 for value in user_symptoms_vect.flatten())95 96 return user_symptoms_vect97 98 99def get_features_fn(*checked_symptoms: Tuple[str]) -> Dict:100 """101 Get vector features based on the selected symptoms.102 103 Args:104 checked_symptoms (Tuple[str]): User symptoms105 106 Returns:107 Dict: The encoded user vector symptoms.108 """109 if not any(lst for lst in checked_symptoms if lst):110 return {111 error_box1: gr.update(visible=True, value="⚠️ Please provide your chief complaints."),112 }113 114 if len(pretty_print(checked_symptoms)) < 5:115 print("Provide at least 5 symptoms.")116 return {117 error_box1: gr.update(visible=True, value="⚠️ Provide at least 5 symptoms"),118 one_hot_vect: None,119 }120 121 return {122 error_box1: gr.update(visible=False),123 one_hot_vect: gr.update(124 visible=False,125 value=get_user_symptoms_from_checkboxgroup(pretty_print(checked_symptoms)),126 ),127 submit_btn: gr.update(value="Data submitted ✅"),128 }129 130 131def key_gen_fn(user_symptoms: List[str]) -> Dict:132 """133 Generate keys for a given user.134 135 Args:136 user_symptoms (List[str]): The vector symptoms provided by the user.137 138 Returns:139 dict: A dictionary containing the generated keys and related information.140 141 """142 clean_directory()143 144 if is_none(user_symptoms):145 print("Error: Please submit your symptoms or select a default disease.")146 return {147 error_box2: gr.update(visible=True, value="⚠️ Please submit your symptoms first."),148 }149 150 # Generate a random user ID151 user_id = np.random.randint(0, 2**32)152 print(f"Your user ID is: {user_id}....")153 154 client = FHEModelClient(path_dir=DEPLOYMENT_DIR, key_dir=KEYS_DIR / f"{user_id}")155 client.load()156 157 # Creates the private and evaluation keys on the client side158 client.generate_private_and_evaluation_keys()159 160 # Get the serialized evaluation keys161 serialized_evaluation_keys = client.get_serialized_evaluation_keys()162 assert isinstance(serialized_evaluation_keys, bytes)163 164 # Save the evaluation key165 evaluation_key_path = KEYS_DIR / f"{user_id}/evaluation_key"166 with evaluation_key_path.open("wb") as f:167 f.write(serialized_evaluation_keys)168 169 serialized_evaluation_keys_shorten_hex = serialized_evaluation_keys.hex()[:INPUT_BROWSER_LIMIT]170 171 return {172 error_box2: gr.update(visible=False),173 key_box: gr.update(visible=False, value=serialized_evaluation_keys_shorten_hex),174 user_id_box: gr.update(visible=False, value=user_id),175 key_len_box: gr.update(176 visible=False, value=f"{len(serialized_evaluation_keys) / (10**6):.2f} MB"177 ),178 gen_key_btn: gr.update(value="Keys have been generated ✅")179 }180 181 182def encrypt_fn(user_symptoms: np.ndarray, user_id: str) -> None:183 """184 Encrypt the user symptoms vector in the `Client Side`.185 186 Args:187 user_symptoms (List[str]): The vector symptoms provided by the user188 user_id (user): The current user's ID189 """190 191 if is_none(user_id) or is_none(user_symptoms):192 print("Error in encryption step: Provide your symptoms and generate the evaluation keys.")193 return {194 error_box3: gr.update(195 visible=True,196 value="⚠️ Please ensure that your symptoms have been submitted and "197 "that you have generated the evaluation key.",198 )199 }200 201 # Retrieve the client API202 client = FHEModelClient(path_dir=DEPLOYMENT_DIR, key_dir=KEYS_DIR / f"{user_id}")203 client.load()204 205 user_symptoms = np.fromstring(user_symptoms[2:-2], dtype=int, sep=".").reshape(1, -1)206 # quant_user_symptoms = client.model.quantize_input(user_symptoms)207 208 encrypted_quantized_user_symptoms = client.quantize_encrypt_serialize(user_symptoms)209 assert isinstance(encrypted_quantized_user_symptoms, bytes)210 encrypted_input_path = KEYS_DIR / f"{user_id}/encrypted_input"211 212 with encrypted_input_path.open("wb") as f:213 f.write(encrypted_quantized_user_symptoms)214 215 encrypted_quantized_user_symptoms_shorten_hex = encrypted_quantized_user_symptoms.hex()[216 :INPUT_BROWSER_LIMIT217 ]218 219 return {220 error_box3: gr.update(visible=False),221 one_hot_vect_box: gr.update(visible=True, value=user_symptoms),222 enc_vect_box: gr.update(visible=True, value=encrypted_quantized_user_symptoms_shorten_hex),223 }224 225 226def send_input_fn(user_id: str, user_symptoms: np.ndarray) -> Dict:227 """Send the encrypted data and the evaluation key to the server.228 229 Args:230 user_id (str): The current user's ID231 user_symptoms (np.ndarray): The user symptoms232 """233 234 if is_none(user_id) or is_none(user_symptoms):235 return {236 error_box4: gr.update(237 visible=True,238 value="⚠️ Please check your connectivity \n"239 "⚠️ Ensure that the symptoms have been submitted and the evaluation "240 "key has been generated before sending the data to the server.",241 )242 }243 244 evaluation_key_path = KEYS_DIR / f"{user_id}/evaluation_key"245 encrypted_input_path = KEYS_DIR / f"{user_id}/encrypted_input"246 247 if not evaluation_key_path.is_file():248 print(249 "Error Encountered While Sending Data to the Server: "250 f"The key has been generated correctly - {evaluation_key_path.is_file()=}"251 )252 253 return {254 error_box4: gr.update(visible=True, value="⚠️ Please generate the private key first.")255 }256 257 if not encrypted_input_path.is_file():258 print(259 "Error Encountered While Sending Data to the Server: The data has not been encrypted "260 f"correctly on the client side - {encrypted_input_path.is_file()=}"261 )262 return {263 error_box4: gr.update(264 visible=True,265 value="⚠️ Please encrypt the data with the private key first.",266 ),267 }268 269 # Define the data and files to post270 data = {271 "user_id": user_id,272 "input": user_symptoms,273 }274 275 files = [276 ("files", open(encrypted_input_path, "rb")),277 ("files", open(evaluation_key_path, "rb")),278 ]279 280 # Send the encrypted input and evaluation key to the server281 url = SERVER_URL + "send_input"282 with requests.post(283 url=url,284 data=data,285 files=files,286 ) as response:287 print(f"Sending Data: {response.ok=}")288 return {289 error_box4: gr.update(visible=False),290 srv_resp_send_data_box: "Data sent",291 }292 293 294def run_fhe_fn(user_id: str) -> Dict:295 """Send the encrypted input and the evaluation key to the server.296 297 Args:298 user_id (int): The current user's ID.299 """300 if is_none(user_id):301 return {302 error_box5: gr.update(303 visible=True,304 value="⚠️ Please check your connectivity \n"305 "⚠️ Ensure that the symptoms have been submitted, the evaluation "306 "key has been generated and the server received the data "307 "before processing the data.",308 ),309 fhe_execution_time_box: None,310 }311 312 data = {313 "user_id": user_id,314 }315 316 url = SERVER_URL + "run_fhe"317 318 with requests.post(319 url=url,320 data=data,321 ) as response:322 if not response.ok:323 return {324 error_box5: gr.update(325 visible=True,326 value=(327 "⚠️ An error occurred on the Server Side. "328 "Please check connectivity and data transmission."329 ),330 ),331 fhe_execution_time_box: gr.update(visible=False),332 }333 else:334 time.sleep(1)335 print(f"response.ok: {response.ok}, {response.json()} - Computed")336 337 return {338 error_box5: gr.update(visible=False),339 fhe_execution_time_box: gr.update(visible=True, value=f"{response.json():.2f} seconds"),340 }341 342 343def get_output_fn(user_id: str, user_symptoms: np.ndarray) -> Dict:344 """Retreive the encrypted data from the server.345 346 Args:347 user_id (str): The current user's ID348 user_symptoms (np.ndarray): The user symptoms349 """350 351 if is_none(user_id) or is_none(user_symptoms):352 return {353 error_box6: gr.update(354 visible=True,355 value="⚠️ Please check your connectivity \n"356 "⚠️ Ensure that the server has successfully processed and transmitted the data to the client.",357 )358 }359 360 data = {361 "user_id": user_id,362 }363 364 # Retrieve the encrypted output365 url = SERVER_URL + "get_output"366 with requests.post(367 url=url,368 data=data,369 ) as response:370 if response.ok:371 print(f"Receive Data: {response.ok=}")372 373 encrypted_output = response.content374 375 # Save the encrypted output to bytes in a file as it is too large to pass through376 # regular Gradio buttons (see https://github.com/gradio-app/gradio/issues/1877)377 encrypted_output_path = CLIENT_DIR / f"{user_id}_encrypted_output"378 379 with encrypted_output_path.open("wb") as f:380 f.write(encrypted_output)381 return {error_box6: gr.update(visible=False), srv_resp_retrieve_data_box: "Data received"}382 383 384def decrypt_fn(385 user_id: str, user_symptoms: np.ndarray, *checked_symptoms, threshold: int = 0.5386) -> Dict:387 """Dencrypt the data on the `Client Side`.388 389 Args:390 user_id (str): The current user's ID391 user_symptoms (np.ndarray): The user symptoms392 threshold (float): Probability confidence threshold393 394 Returns:395 Decrypted output396 """397 398 if is_none(user_id) or is_none(user_symptoms):399 return {400 error_box7: gr.update(401 visible=True,402 value="⚠️ Please check your connectivity \n"403 "⚠️ Ensure that the client has successfully received the data from the server.",404 )405 }406 407 # Get the encrypted output path408 encrypted_output_path = CLIENT_DIR / f"{user_id}_encrypted_output"409 410 if not encrypted_output_path.is_file():411 print("Error in decryption step: Please run the FHE execution, first.")412 return {413 error_box7: gr.update(414 visible=True,415 value="⚠️ Please ensure that: \n"416 "- the connectivity \n"417 "- the symptoms have been submitted \n"418 "- the evaluation key has been generated \n"419 "- the server processed the encrypted data \n"420 "- the Client received the data from the Server before decrypting the prediction",421 ),422 decrypt_box: None,423 }424 425 # Load the encrypted output as bytes426 with encrypted_output_path.open("rb") as f:427 encrypted_output = f.read()428 429 # Retrieve the client API430 client = FHEModelClient(path_dir=DEPLOYMENT_DIR, key_dir=KEYS_DIR / f"{user_id}")431 client.load()432 433 # Deserialize, decrypt and post-process the encrypted output434 output = client.deserialize_decrypt_dequantize(encrypted_output)435 436 top3_diseases = np.argsort(output.flatten())[-3:][::-1]437 top3_proba = output[0][top3_diseases]438 439 out = ""440 441 if top3_proba[0] < threshold or abs(top3_proba[0] - top3_proba[1]) < 0.1:442 out = (443 "⚠️ The prediction appears uncertain; including more symptoms "444 "may improve the results.\n\n"445 )446 447 out = (448 f"{out}Given the symptoms you provided: "449 f"{pretty_print(checked_symptoms, case_conversion=str.capitalize, delimiter=', ')}\n\n"450 "Here are the top3 predictions:\n\n"451 f"1. « {get_disease_name(top3_diseases[0])} » with a probability of {top3_proba[0]:.2%}\n"452 f"2. « {get_disease_name(top3_diseases[1])} » with a probability of {top3_proba[1]:.2%}\n"453 f"3. « {get_disease_name(top3_diseases[2])} » with a probability of {top3_proba[2]:.2%}\n"454 )455 456 return {457 error_box7: gr.update(visible=False),458 decrypt_box: out,459 submit_btn: gr.update(value="Submit"),460 }461 462 463def reset_fn():464 """Reset the space and clear all the box outputs."""465 466 clean_directory()467 468 return {469 one_hot_vect: None,470 one_hot_vect_box: None,471 enc_vect_box: gr.update(visible=True, value=None),472 quant_vect_box: gr.update(visible=False, value=None),473 user_id_box: gr.update(visible=False, value=None),474 default_symptoms: gr.update(visible=True, value=None),475 default_disease_box: gr.update(visible=True, value=None),476 key_box: gr.update(visible=True, value=None),477 key_len_box: gr.update(visible=False, value=None),478 fhe_execution_time_box: gr.update(visible=True, value=None),479 decrypt_box: None,480 submit_btn: gr.update(value="Submit"),481 error_box7: gr.update(visible=False),482 error_box1: gr.update(visible=False),483 error_box2: gr.update(visible=False),484 error_box3: gr.update(visible=False),485 error_box4: gr.update(visible=False),486 error_box5: gr.update(visible=False),487 error_box6: gr.update(visible=False),488 srv_resp_send_data_box: None,489 srv_resp_retrieve_data_box: None,490 **{box: None for box in check_boxes},491 }492 493 494if __name__ == "__main__":495 496 print("Starting demo ...")497 498 clean_directory()499 500 (X_train, X_test), (y_train, y_test), valid_symptoms, diseases = load_data()501 502 with gr.Blocks() as demo:503 504 # Link + images505 gr.Markdown()506 gr.Markdown(507 """508 509 """)510 gr.Markdown()511 gr.Markdown("""<h2 align="center">Health Prediction On Encrypted Data Using Fully Homomorphic Encryption</h2>""")512 gr.Markdown()513 gr.Markdown(514 """515 516 """)517 gr.Markdown()518 gr.Markdown(519 """"520 <p align="center">521 <img width="65%" height="25%" src="https://raw.githubusercontent.com/kcelia/Img/main/healthcare_prediction.jpg">522 </p>523 """524 )525 gr.Markdown("## Notes")526 gr.Markdown(527 """528 - The private key is used to encrypt and decrypt the data and shall never be shared.529 - The evaluation key is a public key that the server needs to process encrypted data.530 """531 )532 533 # ------------------------- Step 1 -------------------------534 gr.Markdown("\n")535 gr.Markdown("## Step 1: Select chief complaints")536 gr.Markdown("<hr />")537 gr.Markdown("<span style='color:grey'>Client Side</span>")538 gr.Markdown("Select at least 5 chief complaints from the list below.")539 540 # Step 1.1: Provide symptoms541 check_boxes = []542 with gr.Row():543 with gr.Column():544 for category in SYMPTOMS_LIST[:3]:545 with gr.Accordion(pretty_print(category.keys()), open=False):546 check_box = gr.CheckboxGroup(pretty_print(category.values()), show_label=0)547 check_boxes.append(check_box)548 with gr.Column():549 for category in SYMPTOMS_LIST[3:6]:550 with gr.Accordion(pretty_print(category.keys()), open=False):551 check_box = gr.CheckboxGroup(pretty_print(category.values()), show_label=0)552 check_boxes.append(check_box)553 with gr.Column():554 for category in SYMPTOMS_LIST[6:]:555 with gr.Accordion(pretty_print(category.keys()), open=False):556 check_box = gr.CheckboxGroup(pretty_print(category.values()), show_label=0)557 check_boxes.append(check_box)558 559 error_box1 = gr.Textbox(label="Error ❌", visible=False)560 561 # Default disease, picked from the dataframe562 gr.Markdown(563 "You can choose an **existing disease** and explore its associated symptoms.",564 visible=False,565 )566 567 with gr.Row():568 with gr.Column(scale=2):569 default_disease_box = gr.Dropdown(sorted(diseases), label="Diseases", visible=False)570 with gr.Column(scale=5):571 default_symptoms = gr.Textbox(label="Related Symptoms:", visible=False)572 # User vector symptoms encoded in oneHot representation573 one_hot_vect = gr.Textbox(visible=False)574 # Submit botton575 submit_btn = gr.Button("Submit")576 # Clear botton577 clear_button = gr.Button("Reset Space 🔁", visible=False)578 579 default_disease_box.change(580 fn=display_default_symptoms_fn, inputs=[default_disease_box], outputs=[default_symptoms]581 )582 583 submit_btn.click(584 fn=get_features_fn,585 inputs=[*check_boxes],586 outputs=[one_hot_vect, error_box1, submit_btn],587 )588 589 # ------------------------- Step 2 -------------------------590 gr.Markdown("\n")591 gr.Markdown("## Step 2: Encrypt data")592 gr.Markdown("<hr />")593 gr.Markdown("<span style='color:grey'>Client Side</span>")594 # Step 2.1: Key generation595 gr.Markdown(596 "### Key Generation\n\n"597 "In FHE schemes, a secret (enc/dec)ryption keys are generated for encrypting and decrypting data owned by the client. \n\n"598 "Additionally, a public evaluation key is generated, enabling external entities to perform homomorphic operations on encrypted data, without the need to decrypt them. \n\n"599 "The evaluation key will be transmitted to the server for further processing."600 )601 602 gen_key_btn = gr.Button("Generate the private and evaluation keys.")603 error_box2 = gr.Textbox(label="Error ❌", visible=False)604 user_id_box = gr.Textbox(label="User ID:", visible=False)605 key_len_box = gr.Textbox(label="Evaluation Key Size:", visible=False)606 key_box = gr.Textbox(label="Evaluation key (truncated):", max_lines=3, visible=False)607 608 gen_key_btn.click(609 key_gen_fn,610 inputs=one_hot_vect,611 outputs=[612 key_box,613 user_id_box,614 key_len_box,615 error_box2,616 gen_key_btn,617 ],618 )619 620 # Step 2.2: Encrypt data locally621 gr.Markdown("### Encrypt the data")622 encrypt_btn = gr.Button("Encrypt the data using the private secret key")623 error_box3 = gr.Textbox(label="Error ❌", visible=False)624 quant_vect_box = gr.Textbox(label="Quantized Vector:", visible=False)625 626 with gr.Row():627 with gr.Column():628 one_hot_vect_box = gr.Textbox(label="User Symptoms Vector:", max_lines=10)629 with gr.Column():630 enc_vect_box = gr.Textbox(label="Encrypted Vector:", max_lines=10)631 632 encrypt_btn.click(633 encrypt_fn,634 inputs=[one_hot_vect, user_id_box],635 outputs=[636 one_hot_vect_box,637 enc_vect_box,638 error_box3,639 ],640 )641 # Step 2.3: Send encrypted data to the server642 gr.Markdown(643 "### Send the encrypted data to the <span style='color:grey'>Server Side</span>"644 )645 error_box4 = gr.Textbox(label="Error ❌", visible=False)646 647 with gr.Row().style(equal_height=False):648 with gr.Column(scale=4):649 send_input_btn = gr.Button("Send data")650 with gr.Column(scale=1):651 srv_resp_send_data_box = gr.Checkbox(label="Data Sent", show_label=False)652 653 send_input_btn.click(654 send_input_fn,655 inputs=[user_id_box, one_hot_vect],656 outputs=[error_box4, srv_resp_send_data_box],657 )658 659 # ------------------------- Step 3 -------------------------660 gr.Markdown("\n")661 gr.Markdown("## Step 3: Run the FHE evaluation")662 gr.Markdown("<hr />")663 gr.Markdown("<span style='color:grey'>Server Side</span>")664 gr.Markdown(665 "Once the server receives the encrypted data, it can process and compute the output without ever decrypting the data just as it would on clear data.\n\n"666 ""667 )668 669 run_fhe_btn = gr.Button("Run the FHE evaluation")670 error_box5 = gr.Textbox(label="Error ❌", visible=False)671 fhe_execution_time_box = gr.Textbox(label="Total FHE Execution Time:", visible=True)672 run_fhe_btn.click(673 run_fhe_fn,674 inputs=[user_id_box],675 outputs=[fhe_execution_time_box, error_box5],676 )677 678 # ------------------------- Step 4 -------------------------679 gr.Markdown("\n")680 gr.Markdown("## Step 4: Decrypt the data")681 gr.Markdown("<hr />")682 gr.Markdown("<span style='color:grey'>Client Side</span>")683 gr.Markdown(684 "### Get the encrypted data from the <span style='color:grey'>Server Side</span>"685 )686 687 error_box6 = gr.Textbox(label="Error ❌", visible=False)688 689 # Step 4.1: Data transmission690 with gr.Row().style(equal_height=True):691 with gr.Column(scale=4):692 get_output_btn = gr.Button("Get data")693 with gr.Column(scale=1):694 srv_resp_retrieve_data_box = gr.Checkbox(label="Data Received", show_label=False)695 696 get_output_btn.click(697 get_output_fn,698 inputs=[user_id_box, one_hot_vect],699 outputs=[srv_resp_retrieve_data_box, error_box6],700 )701 702 # Step 4.1: Data transmission703 gr.Markdown("### Decrypt the output")704 decrypt_btn = gr.Button("Decrypt the output using the private secret key")705 error_box7 = gr.Textbox(label="Error ❌", visible=False)706 decrypt_box = gr.Textbox(label="Decrypted Output:")707 708 decrypt_btn.click(709 decrypt_fn,710 inputs=[user_id_box, one_hot_vect, *check_boxes],711 outputs=[decrypt_box, error_box7, submit_btn],712 )713 714 # ------------------------- End -------------------------715 716 gr.Markdown(717 """T718 """719 )720 721 gr.Markdown("\n\n")722 723 gr.Markdown(724 """"""725 )726 727 clear_button.click(728 reset_fn,729 outputs=[730 one_hot_vect_box,731 one_hot_vect,732 submit_btn,733 error_box1,734 error_box2,735 error_box3,736 error_box4,737 error_box5,738 error_box6,739 error_box7,740 default_disease_box,741 default_symptoms,742 user_id_box,743 key_len_box,744 key_box,745 quant_vect_box,746 enc_vect_box,747 srv_resp_send_data_box,748 srv_resp_retrieve_data_box,749 fhe_execution_time_box,750 decrypt_box,751 *check_boxes,752 ],753 )754 755 demo.launch()756