chore: add relevent files

app.py

@@ -0,0 +1,1014 @@
+import copy
+import io
+import random
+import subprocess
+import threading
+import time
+from glob import glob
+from pathlib import Path
+
+import gradio as gr
+import matplotlib.pyplot as plt
+import numpy
+import requests
+from PIL import Image
+from scipy import stats
+from server import *
+from tqdm import tqdm
+from utils import *
+
+from concrete.ml.deployment import FHEModelClient
+
+CURRENT_DIR = Path(__file__).parent
+
+subprocess.Popen(["uvicorn", "server:app"], cwd=CURRENT_DIR)
+time.sleep(3)
+
+numpy.set_printoptions(threshold=numpy.inf)
+
+USER_ID = numpy.random.randint(0, 2**32)
+
+# Define client-specific directories
+CLIENT_DIR = ROOT_DIR / f"user_{USER_ID}/client"
+CLIENT_KEY_SMOOTHER_MODULE_DIR = CLIENT_DIR / KEY_SMOOTHER_MODULE_DIR
+CLIENT_KEY_BASE_MODULE_DIR = CLIENT_DIR / KEY_BASE_MODULE_DIR
+CLIENT_ENCRYPTED_INPUT_DIR = CLIENT_DIR / ENCRYPTED_INPUT_DIR
+CLIENT_ENCRYPTED_OUTPUT_DIR = CLIENT_DIR / ENCRYPTED_OUTPUT_DIR
+
+# Define server-specific directories
+SERVER_DIR = ROOT_DIR / f"user_{USER_ID}/server"
+SERVER_KEY_SMOOTHER_MODULE_DIR = SERVER_DIR / KEY_SMOOTHER_MODULE_DIR
+SERVER_KEY_BASE_MODULE_DIR = SERVER_DIR / KEY_BASE_MODULE_DIR
+SERVER_ENCRYPTED_INPUT_DIR = SERVER_DIR / ENCRYPTED_INPUT_DIR
+SERVER_ENCRYPTED_OUTPUT_DIR = SERVER_DIR / ENCRYPTED_OUTPUT_DIR
+
+ALL_DIRECTORIES = [
+    CLIENT_KEY_SMOOTHER_MODULE_DIR,
+    CLIENT_KEY_BASE_MODULE_DIR,
+    CLIENT_ENCRYPTED_INPUT_DIR,
+    CLIENT_ENCRYPTED_OUTPUT_DIR,
+    SERVER_KEY_SMOOTHER_MODULE_DIR,
+    SERVER_KEY_BASE_MODULE_DIR,
+    SERVER_ENCRYPTED_INPUT_DIR,
+    SERVER_ENCRYPTED_OUTPUT_DIR,
+]
+
+# Load test dataset
+
+print("Load data ...")
+UNIQUE_FOUNDERS = load_pickle_from_zip("./data/unique_founders.pkl")
+OTHER_TEST_FOUNDERS = load_pickle_from_zip("./data/hf_test_founders.pkl")
+OTHER_TRAIN_FOUNDERS = load_pickle_from_zip("./data/hf_train_founders.pkl")
+UNIQUE_MIXED_FOUNDERS = load_pickle_from_zip("./data/unique_mixed_founders.pkl")
+
+DESCENDANT_PATH = Path("./data/Child.pkl")
+PREDICTION_IMG_PATH = Path("output.npg")
+FAMILY_TREE_IMG_PATH = Path("simulated_family_tree.png")
+
+ALL_GENERATED_PATHS = [
+    DESCENDANT_PATH,
+    FAMILY_TREE_IMG_PATH,
+    PREDICTION_IMG_PATH,
+    FHE_COMPUTATION_TIMELINE,
+]
+
+
+def reset():
+    """Reset the environment.
+
+    Clean the root directory, recreating necessary directories and removing any generated files.
+    """
+    print("Cleaning ...")
+
+    clean_dir(ROOT_DIR)
+
+    for directory in ALL_DIRECTORIES:
+        directory.mkdir(parents=True, exist_ok=True)
+
+    for file_path in ALL_GENERATED_PATHS:
+        if file_path.exists():
+            file_path.unlink()
+            print(f"File: {file_path} has been removed.")
+
+
+def simulate_allele_fn():
+
+    yield {
+        simulate_btn: gr.update(visible=True, value="🔄 Processing... Please wait."),
+        ethnicity_simulation_img: gr.update(visible=False),
+        simulate_text: gr.update(visible=False),
+    }
+
+    start_time = time.time()
+    n_generations = random.randint(1, 3)
+
+    individuals = random.sample(UNIQUE_MIXED_FOUNDERS, 2 + n_generations)
+
+    num_snps = META["C"]
+    print(f"Simulating family tree with: {n_generations} generations ...")
+
+    first_founder = individuals.pop(0)
+    second_founder = individuals.pop(0)
+
+    assert len(first_founder) == 4
+    assert len(second_founder) == 4
+    assert (
+        sum([numpy.array_equal(arr1, arr2) for arr1, arr2 in zip(first_founder, second_founder)])
+        == 0
+    )
+
+    lineages = []
+    labels = []
+    admix = []
+
+    for gen in range(n_generations + 1):
+
+        print(f"Generation_{gen}:")
+
+        # Initialize the child for this generation
+        if gen == 0:
+            # Use the specified founders for the first generation
+            founder_1 = first_founder
+            founder_2 = second_founder
+            labels.append(["Ancestor 1", "Ancestor 2", "Progeny (Generation 1)"])
+        else:
+            # Use the last child from the previous generation as founder 1
+            founder_1 = admix[-1]
+            founder_2 = individuals.pop(0)
+            labels.append(
+                [labels[-1][-1], f"Ancestor {len(labels) + 2}", f"Progeny (Generation {gen + 1})"]
+            )
+
+        assert len(founder_1) == 4
+        assert len(founder_2) == 4
+
+        # Prepare new admix entry for this generation's child
+        admix.append([None, None, None, None])
+
+        snp_1, snp_2, label_1, label_2 = copy.deepcopy(founder_1)
+        _snp_1, _snp_2, _label_1, _label_2 = copy.deepcopy(founder_2)
+
+        lineage = []
+        for j in range(2):  # Two haplotypes
+
+            # Select one haplotype from each founder
+            snp, label = (snp_1, label_1) if random.random() < 0.5 else (snp_2, label_2)
+            _snp, _label = (_snp_1, _label_1) if random.random() < 0.5 else (_snp_2, _label_2)
+
+            lineage.append(
+                [compute_distribution(label.flatten()), compute_distribution(_label.flatten())]
+            )
+
+            breakpoints = numpy.random.choice(
+                range(1, num_snps),
+                # size=int(sum(numpy.random.poisson(0.75, size=gen))) + 1,
+                size=int(sum(numpy.random.poisson(0.1, size=gen)))
+                + 1,  # Fewer breakpoints, less mixed
+                replace=False,
+            )
+
+            breakpoints = numpy.concatenate(([0], numpy.sort(breakpoints), [num_snps]))
+
+            for k in range(len(breakpoints) - 1):
+
+                snp[breakpoints[k] : breakpoints[k + 1]] = _snp[
+                    breakpoints[k] : breakpoints[k + 1]
+                ].copy()
+                label[breakpoints[k] : breakpoints[k + 1]] = _label[
+                    breakpoints[k] : breakpoints[k + 1]
+                ].copy()
+
+                yield {
+                    simulate_btn: gr.update(
+                        visible=True,
+                        value=f"⏳ Time elapsed: {time.time() - start_time:.0f} seconds ({(gen + 1)/(n_generations + 1):.0%})",
+                    )
+                }
+
+            admix[-1][j], admix[-1][j + 2] = snp, label
+
+        lineages.append([(lineage[0][0] + lineage[1][0]) / 2, (lineage[0][1] + lineage[1][1]) / 2])
+
+        print(f"Ascendant_1: {lineages[-1][0]} + Ascendant_2 {lineages[-1][0]}")
+
+    last_child = admix[-1]
+    snp_1, snp_2, label_1, label_2 = last_child[0], last_child[1], last_child[2], last_child[3]
+    snp, label_full = (snp_1, label_1) if random.random() < 0.5 else (snp_2, label_2)
+
+    # We are prediction on one allele, so we plot the right allele
+    lineages[-1][-1] = compute_distribution(label_full)
+
+    l1 = label_full.reshape(1, -1)
+    N, L = l1.shape
+    y = l1[:, 0 : L // META["M"] * META["M"]].reshape(N, L // META["M"], META["M"])
+    y = stats.mode(y, axis=2)[0].squeeze()
+
+    write_pickle(path="./data/Child.pkl", data=[snp.reshape(1, -1), y])
+    snp = numpy.array(snp).reshape(1, -1)
+    yield {
+        simulate_btn: gr.update(
+            visible=True,
+            value=f"⏳ Time elapsed: {time.time() - start_time:.0f} seconds ({(gen + 1)/(n_generations + 1):.0%})",
+        )
+    }
+
+    print("Plot the simulated allele.")
+    print(f"{snp.shape=} - {y.shape=}")
+    print(
+        f"{any(snp.flatten()[12343 : 12343 + 1000] == 1)=} - {any(snp.flatten()[12343 : 12343 + 1000] == 0)=}"
+    )
+
+    yield {
+        simulate_btn: gr.update(
+            visible=True,
+            value=f"⏳ Time elapsed: {time.time() - start_time:.0f} seconds ({(gen + 1)/(n_generations + 1):.0%})",
+        )
+    }
+
+    _ = pie_ethnicity_simulation_plot_img(copy.copy(lineages), copy.copy(labels))
+
+    sorted_indices = numpy.argsort(lineages[-1][0])[::-1][:2]
+    top_percentages = [lineages[-1][0][i] for i in sorted_indices]
+    top_labels = [LABELS[i] for i in sorted_indices]
+
+    # items = [f'{p:.0%} {l}' for p, l in zip(top_percentages, top_labels)]
+    # items = [f'{p:.0%} {l}' for p, l in zip(lineages[-1][0], LABELS)]
+
+    yield {
+        simulate_btn: gr.update(
+            visible=True,
+            value=f"⏳ Time elapsed: {time.time() - start_time:.0f} seconds ({(gen + 1)/(n_generations + 1):.0%})",
+        )
+    }
+
+    yield {
+        clear_input_box: gr.update(visible=True, value=list(snp.flatten())[:321]),
+        simulate_btn: gr.update(value="Data simulated ✅"),
+        ethnicity_simulation_img: gr.update(
+            value=Image.open(FAMILY_TREE_IMG_PATH),
+            visible=True,
+            show_label=False,
+            show_download_button=True,
+            container=True,
+        ),
+        simulate_text: gr.update(
+            value=f"Given the genetic lineage simulation above, the origin of two predominant genes, for the last progeny are: {top_labels[-2]} and {top_labels[-1]}. Now, we proceed with ***DNA testing*** using **FHE** on this final descendant.",
+            visible=True,
+        ),
+    }
+
+    return
+
+
+def key_gen_fn(user_id):
+    """Generate keys for a given user on the Client Side."""
+
+    print("\n------------ Step 1: Key Generation:")
+
+    yield gr.update(visible=True, value="🔄 Processing... Please wait.")
+
+    print(f"Your user ID is: {user_id:.0f}....")
+
+    ## Generate one key for all models since they share the same crypto params
+    # Stage1: Base modules
+
+    base_modules_path = sorted(glob(f"{SHARED_BASE_MODULE_DIR}/model_*"), key=extract_model_number)
+
+    print(f"{len(base_modules_path)=} {META['NW']=}")
+
+    if len(base_modules_path) != META["NW"]:
+        yield gr.update(visible=True, value="❌ Error in key generation", interactive=False)
+
+    base_client = FHEModelClient(path_dir=base_modules_path[0], key_dir=CLIENT_KEY_BASE_MODULE_DIR)
+    base_client.generate_private_and_evaluation_keys()
+    serialized_evaluation_base_modules_keys = base_client.get_serialized_evaluation_keys()
+    assert isinstance(serialized_evaluation_base_modules_keys, bytes)
+    print(f"Stage1: {len(glob(f'{CLIENT_KEY_BASE_MODULE_DIR}/eval_key'))} key has been generated")
+
+    # Stage2: Smoother module
+    smoother_client = FHEModelClient(
+        path_dir=SHARED_SMOOTHER_MODULE_DIR, key_dir=CLIENT_KEY_SMOOTHER_MODULE_DIR
+    )
+    smoother_client.generate_private_and_evaluation_keys()
+    serialized_evaluation_smoother_module_keys = smoother_client.get_serialized_evaluation_keys()
+    assert isinstance(serialized_evaluation_smoother_module_keys, bytes)
+    print(
+        f"Stage2: {len(glob(f'{CLIENT_KEY_SMOOTHER_MODULE_DIR}/eval_key'))} key has been generated"
+    )
+
+    # Save the keys
+    base_evaluation_key_path = Path(base_client.key_dir) / "eval_key"
+    smoother_evaluation_key_path = Path(smoother_client.key_dir) / "eval_key"
+
+    write_bytes(base_evaluation_key_path, serialized_evaluation_base_modules_keys)
+    write_bytes(smoother_evaluation_key_path, serialized_evaluation_smoother_module_keys)
+
+    if not base_evaluation_key_path.is_file():
+        msg = "❌ Error encountered while generating the base modules key evaluation"
+    elif not smoother_evaluation_key_path.is_file():
+        msg = "❌ Error encountered while generating the smoother module key evaluation"
+    else:
+        msg = "Secret and public keys have been generated ✅"
+
+    print(msg)
+
+    yield gr.update(visible=True, value=msg, interactive=False)
+    return
+
+
+def encrypt_fn(user_id):
+    """Encrypt input on the Client Side using the secret key."""
+
+    print("\n------------ Step 2: Encrypt the input")
+
+    if (
+        is_none(int(user_id))
+        or (len(glob(f"{CLIENT_KEY_BASE_MODULE_DIR}/*")) == 0)
+        or not DESCENDANT_PATH.is_file()
+    ):
+
+        print("Error in encryption step: Provide your chromosome and generate the evaluation keys.")
+
+        yield {
+            encrypt_btn: gr.update(
+                visible=True,
+                value="❌ Ensure your have simulated an allele and the evaluation key has been generated.",
+            )
+        }
+        return
+
+    allele, _ = read_pickle(path=DESCENDANT_PATH)
+
+    yield {
+        encrypt_btn: gr.update(visible=True, value="🔄 Processing... Please wait."),
+        send_btn: gr.update(interactive=False),
+        run_fhe_btn: gr.update(interactive=False),
+        get_output_btn: gr.update(interactive=False),
+        decrypt_btn: gr.update(interactive=False),
+        simulate_btn: gr.update(interactive=False),
+    }
+
+    base_modules_path = sorted(glob(f"{SHARED_BASE_MODULE_DIR}/model_*"), key=extract_model_number)
+    assert len(base_modules_path) == META["NW"]
+
+    print(f"{len(base_modules_path)} models have been loaded")
+
+    client_fhemodels = []
+    for i, base_module_path in enumerate(tqdm(base_modules_path)):
+        base_client = FHEModelClient(path_dir=base_module_path, key_dir=CLIENT_KEY_BASE_MODULE_DIR)
+        client_fhemodels.append(base_client)
+    base_serialized_evaluation_keys = read_bytes(base_client.key_dir / "eval_key")
+    assert isinstance(base_serialized_evaluation_keys, bytes)
+
+    X_p, _, M_, rem = process_data_for_base_modules(META, allele)
+    base_args = tuple(zip(client_fhemodels[:-1], numpy.swapaxes(X_p, 0, 1)))
+    base_args += ((client_fhemodels[-1], allele[:, allele.shape[1] - (M_ + rem) :]),)
+
+    start_time = time.time()
+    for i, (client, window) in enumerate(base_args):
+        encrypted_input = client.quantize_encrypt_serialize(window)
+        write_bytes(CLIENT_ENCRYPTED_INPUT_DIR / f"window_{i}", encrypted_input)
+        yield {
+            encrypt_btn: gr.update(
+                visible=True,
+                value=f"⏳ Time elapsed: {time.time() - start_time:.0f} seconds ({i/META['NW']:.0%}).",
+            )
+        }
+        # f"⏳ Time elapsed: {time.time() - start_time:.0f} seconds ({(gen + 1)/(n_generations + 1):.0%})")
+    exec_time = time.time() - start_time
+
+    msg = f"Encryption completed in {exec_time: .2f} seconds."
+
+    print(msg)
+
+    enc_quant_input_shorten_hex = encrypted_input.hex()[:INPUT_BROWSER_LIMIT]
+
+    yield {
+        encrypt_input_box: gr.update(visible=True, value=enc_quant_input_shorten_hex),
+        encrypt_btn: gr.update(interactive=False, value=msg),
+        simulate_btn: gr.update(interactive=False),
+        send_btn: gr.update(interactive=True),
+        run_fhe_btn: gr.update(interactive=True),
+        get_output_btn: gr.update(interactive=True),
+        decrypt_btn: gr.update(interactive=True),
+    }
+    return
+
+
+def send_input_fn(user_id):
+    """Send the encrypted data and the evaluation key to the server."""
+
+    print("\n------------ Step 3.1: Send encrypted_data to the Server")
+
+    errors = []
+    if not (CLIENT_KEY_BASE_MODULE_DIR / "eval_key").is_file():
+        errors.append("Stage 1 evaluation keys are missing.")
+
+    if not (CLIENT_KEY_SMOOTHER_MODULE_DIR / "eval_key").is_file():
+        errors.append("Stage 2 evaluation keys are missing. ")
+
+    if len(glob(str(CLIENT_ENCRYPTED_INPUT_DIR / "window_*"))) != META["NW"]:
+        errors.append("The input has not been successfully encrypted.")
+
+    if errors:
+        error_message = "❌ Error during data transmission:\n" + "\n".join(errors)
+        print(error_message)
+
+        yield {send_btn: gr.update(value=error_message)}
+        return
+
+    yield {
+        send_btn: gr.update(value="🔄 Processing... Please wait."),
+        run_fhe_btn: gr.update(interactive=False),
+        get_output_btn: gr.update(interactive=False),
+        decrypt_btn: gr.update(interactive=False),
+    }
+
+    # Define the data and files to post
+    data = {"user_id": f"user_{user_id:.0f}", "root_dir": str(ROOT_DIR)}
+    n_w = glob(f"{CLIENT_ENCRYPTED_INPUT_DIR}/window_*")
+
+    files = [
+        ("files", open(f"{CLIENT_KEY_BASE_MODULE_DIR}/eval_key", "rb")),
+        ("files", open(f"{CLIENT_KEY_SMOOTHER_MODULE_DIR}/eval_key", "rb")),
+    ] + [("files", open(f"{CLIENT_ENCRYPTED_INPUT_DIR}/window_{i}", "rb")) for i in range(len(n_w))]
+
+    # Send the encrypted input and evaluation key to the server
+    url = SERVER_URL + "send_input"
+    print(f"{url=}")
+
+    with requests.post(url=url, data=data, files=files) as resp:
+        print(f"{resp.ok=}")
+        msg = "Data sent to the Server ✅" if resp.ok else "❌ Error in sending data to the Server"
+
+    yield {
+        send_btn: gr.update(value=msg, interactive=False if "✅" in msg else True),
+        run_fhe_btn: gr.update(interactive=True),
+        get_output_btn: gr.update(interactive=True),
+        decrypt_btn: gr.update(interactive=True),
+    }
+
+    return
+
+
+def run_fhe_fn(user_id):
+    """Run the FHE execution on the Server Side."""
+
+    print("\n------------ Step 4.1: Run in FHE on the Server Side")
+
+    if FHE_COMPUTATION_TIMELINE.exists():
+        FHE_COMPUTATION_TIMELINE.unlink()
+        print(f"File {FHE_COMPUTATION_TIMELINE} removed successfully.")
+
+
+    if is_none(int(user_id)) or len(glob(f"{SERVER_ENCRYPTED_INPUT_DIR}/encrypted_window_*")) == 0:
+        yield {
+            run_fhe_btn: gr.update(
+                visible=True,
+                value="❌ Check your connectivity. Ensure the input has been submitted, the keys have been generated, and the server has received the data.",
+            )
+        }
+        return
+
+    yield {
+        run_fhe_btn: gr.update(
+            visible=True, value="🔄 Processing... Please wait. This may take up to 500 seconds."
+        ),
+        get_output_btn: gr.update(interactive=False),
+        decrypt_btn: gr.update(interactive=False),
+    }
+
+    data = {
+        "user_id": f"user_{user_id:.0f}",
+        "root_dir": str(ROOT_DIR),
+    }
+
+    url = SERVER_URL + "run_fhe"
+
+    # Function to run FHE on the server in a separate thread
+    def run_fhe_on_server():
+        nonlocal server_response
+        with requests.post(url=url, data=data) as resp:
+            if not resp.ok:
+                server_response = "error"
+            else:
+                server_response = resp.json()
+
+    server_response = None
+
+    # Start the FHE process in a separate thread
+    server_thread = threading.Thread(target=run_fhe_on_server)
+    server_thread.start()
+
+    # While the server is processing, check the timing file for updates
+    while server_thread.is_alive():
+        try:
+            with FHE_COMPUTATION_TIMELINE.open("r", encoding="utf-8") as f:
+                timing = f.read().strip()
+            yield {
+                run_fhe_btn: gr.update(visible=True, value=f"⏳ Time elapsed: {timing}"),
+            }
+        except FileNotFoundError:
+            yield {
+                run_fhe_btn: gr.update(visible=True, value="⏳ Waiting for the server to start..."),
+            }
+
+        time.sleep(5)  # Wait a few seconds before reading again
+
+    # Wait for the thread to finish
+    server_thread.join()
+
+    # Handle server response after completion
+    if server_response == "error":
+        yield {
+            run_fhe_btn: gr.update(
+                visible=True,
+                value="❌ Error occurred on the Server Side. Please check your connectivity.",
+            ),
+        }
+    else:
+        final_time = server_response
+        yield {
+            run_fhe_btn: gr.update(
+                visible=True, interactive=False, value=f"FHE executed in {final_time:.2f} seconds"
+            ),
+            get_output_btn: gr.update(interactive=True),
+            decrypt_btn: gr.update(interactive=True),
+        }
+
+
+def get_output_fn(user_id):
+    """Retreive the encrypted data from the server."""
+
+    print("\n------------ Step 5.1: Get output")
+
+    if is_none(int(user_id)) or len(glob(f"{SERVER_ENCRYPTED_INPUT_DIR}/encrypted_window_*")) == 0:
+        msg = "❌ Error during data transmission: The server did not receive the data, so the FHE process could not be performed."
+        print(msg)
+        yield {get_output_btn: gr.update(visible=True, value=msg)}
+        return
+
+    yield {
+        get_output_btn: gr.update(value="🔄 Processing... Please wait."),
+        decrypt_btn: gr.update(interactive=False),
+    }
+
+    data = {
+        "user_id": f"user_{user_id:.0f}",
+        "root_dir": str(ROOT_DIR),
+    }
+
+    # Retrieve the encrypted output
+    url = SERVER_URL + "get_output"
+    print(f"{url=}")
+
+    with requests.post(url=url, data=data) as response:
+        if response.ok:
+            msg = (
+                "Data sent to the Client ✅"
+                if response.ok
+                else "❌ Error in receiving data from the server"
+            )
+            yhat_encrypted = load_pickle_from_zip(CLIENT_ENCRYPTED_OUTPUT_DIR / "encrypted_final_output.pkl")
+            assert len(yhat_encrypted) == META["NW"]
+
+    yield {
+        get_output_btn: gr.update(value=msg, interactive=False if "✅" in msg else True),
+        decrypt_btn: gr.update(interactive=True),
+    }
+    return
+
+
+def decrypt_fn(user_id):
+    """Dencrypt the data on the Client Side."""
+
+    print("\n------------ Step 6: Decrypt output")
+
+    if (
+        is_none(int(user_id))
+        or not (CLIENT_ENCRYPTED_OUTPUT_DIR / "encrypted_final_output.pkl").is_file()
+    ):
+        print("Error in decryption step: Please run the FHE execution, first.")
+        yield {
+            decrypt_btn: gr.update(
+                visible=True,
+                value="❌ Ensure the input is precessed and retrieved from the server",
+            ),
+        }
+
+        return
+
+    yield {decrypt_btn: gr.update(visible=True, value="🔄 Processing... Please wait.")}
+
+    yhat_encrypted = load_pickle_from_zip(CLIENT_ENCRYPTED_OUTPUT_DIR / "encrypted_final_output.pkl")
+
+    # Retrieve the client API
+    client = FHEModelClient(path_dir=SHARED_SMOOTHER_MODULE_DIR, key_dir=SHARED_SMOOTHER_MODULE_DIR)
+    client.load()
+
+    yhat = []
+    for encrypted_i in yhat_encrypted:
+        # Deserialize, decrypt and post-process the encrypted output
+        output = client.deserialize_decrypt_dequantize(encrypted_i)
+        y_pred = numpy.argmax(output, axis=-1)[0]
+        yhat.append(y_pred)
+
+    yhat = numpy.array(yhat)
+
+    proportion = compute_distribution(yhat.flatten())
+    _ = pie_output_plot_img(copy.copy(proportion))
+
+    yield {
+        decrypt_btn: gr.update(value="Output decrypted ✅", interactive=False),
+        pie_plot_output: gr.update(
+            value=Image.open(PREDICTION_IMG_PATH),
+            visible=True,
+            show_label=False,
+            show_download_button=False,
+            container=False,
+        ),
+        user_id_btn: gr.update(value=None),
+    }
+    return
+
+
+def create_pie_chart(ax, data, title, highlight=False, largest_piece=False, simulation=True):
+
+    sorted_indices = numpy.argsort(data)[::-1]
+    if data[sorted_indices[0]] == 1:
+        sorted_indices = [sorted_indices[0]]
+
+    sorted_data = data[sorted_indices]
+    sorted_labels = [LABELS[i] for i in sorted_indices]
+    sorted_colors = [COLORS[i] for i in sorted_indices]
+
+    ## Keep only the 2 biggest parts
+    if simulation and not data[sorted_indices[0]] == 1:
+        top_data = sorted_data[:2]  # First two largest proportions
+        others_data = sorted_data[2:].sum()  # Sum of the rest
+        sorted_data = numpy.concatenate([top_data, [others_data]])  # Include "others"
+
+        sorted_labels = sorted_labels[:2] + ["Others"]  # First two labels + "Others"
+        sorted_colors = sorted_colors[:2] + ["#D3D3D3"]  # Gray color for "Others"
+
+    if highlight:
+        explode = [0.15 for _ in range(len(sorted_data))]
+    else:
+        explode = [0.09 if i == 0 else 0 for i in range(len(sorted_data))]
+
+    wedges, _, _ = ax.pie(
+        sorted_data,
+        labels=sorted_labels,
+        colors=sorted_colors,
+        autopct=lambda x: f"{round(x)}%",
+        pctdistance=0.7,
+        labeldistance=1.1,
+        shadow=True,
+        explode=explode,
+        radius=1.9 if highlight else 1.0,
+    )
+
+    if largest_piece:
+        # Highlight the largest wedge with a black edge
+        wedges[0].set_edgecolor("black")
+        wedges[0].set_linewidth(2)
+
+    ax.set_title(title, fontsize=10, weight="bold")
+
+    if highlight:
+        for wedge in wedges:
+            wedge.set_edgecolor("black")
+            wedge.set_linewidth(3)
+
+        ax.set_title(title, fontsize=14, weight="bold")
+
+    ax.axis("equal")  # Ensure the pie chart is drawn as a circle
+
+
+def pie_ethnicity_simulation_plot_img(lineages, labels):
+    """Generates a pie chart for genetic lineage simulation across multiple generations."""
+
+    n_generations = len(lineages) - 1
+
+    fig, axes = plt.subplots(n_generations, 3, figsize=(10, 4 * n_generations))
+
+    fig.suptitle("Genetic Lineage Simulator", fontsize=16, weight="bold", x=1)
+
+    for gen in range(n_generations):
+        parent1, parent2 = lineages.pop(0)
+        descendant = lineages[0][0]
+        label = labels.pop(0)
+
+        print(f"Generation {gen}: Parent 1: {parent1} + Parent 2 {parent2} = Child {descendant}")
+
+        ax_gen = axes[gen] if n_generations > 1 else axes
+
+        create_pie_chart(ax_gen[0], parent1, label[0])
+        create_pie_chart(ax_gen[1], parent2, label[1])
+
+        # Check if it's the last descendant, highlight it
+        is_last = gen == n_generations - 1
+        create_pie_chart(
+            ax_gen[2],
+            descendant,
+            f"Last progeny (Generation {n_generations + 1})" if is_last else label[2],
+            highlight=is_last,
+        )
+
+    plt.subplots_adjust(right=2)
+
+    plt.savefig(FAMILY_TREE_IMG_PATH, format="png", bbox_inches="tight")
+
+    buf = io.BytesIO()
+    plt.savefig(buf, format="png")
+    buf.seek(0)
+    plt.close(fig)
+
+    return Image.open(buf)
+
+
+def pie_output_plot_img(data, simulation=False):
+    """Generates a pie chart based on the ethnic proportions."""
+
+    fig, ax = plt.subplots(figsize=(10, 8))
+
+    create_pie_chart(
+        ax,
+        data,
+        "Predicted ethnicity distribution using FHE",
+        highlight=True,
+        largest_piece=True,
+        simulation=simulation,
+    )
+
+    plt.savefig(PREDICTION_IMG_PATH, format="png", bbox_inches="tight")
+
+    buf = io.BytesIO()
+    plt.savefig(buf, format="png")
+    buf.seek(0)
+    plt.close(fig)  # Close the plot to free memory
+    return Image.open(buf)
+
+
+CSS = """                
+        #accordion-label { /* Custom styling for the Accordion title */
+            background-color: #f0f0f0  !important; /* Set the background color to gray */
+        }
+
+        
+        #ie_plot_output {  /* Target the image output container */
+            align-items: center;
+            justify-content: center;
+            margin: auto;  /* Ensure it is centered */
+        }
+"""
+
+if __name__ == "__main__":
+
+    print("Starting demo ...")
+
+    print(META)
+
+    reset()
+
+    with gr.Blocks(css=CSS) as demo:
+
+        gr.Markdown()
+        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()
+        gr.Markdown(
+            """<h2 align="center">Encrypted DNA Testing Using Fully Homomorphic Encryption</h2>"""
+        )
+        gr.Markdown()
+        gr.Markdown(
+            """
+            <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="https://user-images.githubusercontent.com/5758427/197972109-faaaff3e-10e2-4ab6-80f5-7531f7cfb08f.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="https://user-images.githubusercontent.com/5758427/197976802-fddd34c5-f59a-48d0-9bff-7ad1b00cb1fb.png">Documentation</a>
+                —
+                <a href="https://zama.ai/community"> <img style="vertical-align: middle; display:inline-block; margin-right: 3px;" width=15 src="https://user-images.githubusercontent.com/5758427/197977153-8c9c01a7-451a-4993-8e10-5a6ed5343d02.png">Community</a>
+                —
+                <a href="https://twitter.com/zama_fhe"> <img style="vertical-align: middle; display:inline-block; margin-right: 3px;" width=15 src="https://user-images.githubusercontent.com/5758427/197975044-bab9d199-e120-433b-b3be-abd73b211a54.png">@zama_fhe</a>
+            </p>
+            """
+        )
+
+        gr.Markdown()
+
+        gr.Markdown(
+            """
+            <p align='center'> DNA testing platforms analyze your genetic data in the clear, leaving it vulnerable to hacks. With Fully Homomorphic Encryption (FHE), they could perform this analysis on encrypted data, ensuring your sensitive information remains safe, even during processing, allowing to get the knowledge without the risks.</p>
+
+            <p align='center'> In this demo, we show you how to perform encrypted DNA analysis using FHE and Zama's Concrete ML <a href='https://github.com/zama-ai/concrete-ml' target=library</a>library</p>
+            """
+        )
+
+        gr.Markdown()
+        gr.Markdown(
+            """
+            <p align="center">
+                <img width="75%" height="30%" src="https://raw.githubusercontent.com/kcelia/Img/main/dna_banner.png">
+            </p>
+            """
+        )
+
+        with gr.Accordion("How does it work?", elem_id="accordion-label", open=False):
+            gr.Markdown(
+                """
+                **In FHE, two special keys are created to make secure, private DNA analysis possible:**
+
+                - Secret key (your personal lock and key): This is like your own unique lock and key. You use it to encrypt (lock up) your DNA
+                data so that no one else can read it. You also use it to decrypt (unlock) the results after they've been processed.
+                - Evaluation key (a safe tool for the DNA company): This is a special key you provide to the DNA analysis company. It allows
+                them to perform their analyses on your encrypted DNA data without decrypting it. Think of it as giving them specialized tools
+                that let them work on your locked box without needing to open it.
+
+                **How it works:**
+
+                1. You encrypt your DNA data: Using your secret key, you lock up your DNA information so it's unreadable to anyone else.
+                2. You send the encrypted DNA data and evaluation key to the DNA company: The company receives your encrypted DNA data
+                along with the evaluation key. They can't see your actual genetic information because it's securely locked.
+                3. The DNA company processes data without seeing it: Using the evaluation key, the company runs their genetic analysis
+                algorithms on your encrypted data. They might assess health risks, ancestry information, or genetic traits—all without ever
+                unlocking your DNA data.
+                4. You decrypt the results: After processing, the company sends back the encrypted results of their analysis. You use your
+                secret key to decrypt and view the findings.
+
+                **In simple terms:**
+
+                FHE allows you to get your DNA analyzed without the DNA company ever seeing your actual genetic code. You keep your DNA
+                data securely locked with your secret key and give the company a special tool (the evaluation key) that lets them perform the
+                analysis without unlocking your data. This way, you receive personalized genetic insights while keeping your most sensitive
+                information private.
+                """
+            )
+
+        gr.Markdown("<hr />")
+        gr.Markdown("# Step 1: Generate Genetic Lineage Simulator")
+
+        gr.Markdown(
+            """
+            To start this demo, the simulation randomly picks *N* individuals over *G* generations, from a genetic dataset.
+
+            Each individual is represented by two alleles from chromosome *22*, which is particularly relevent for tracing human evolutionary history and migration patterns across the world. By analyzing specific markers known as Single Nucleotide Polymorphisms (SNPs) at key positions on this chromosome, valuable insights can be gained.
+
+            Each ancestor transmits half of their genetic material, with *50%* of our genes coming from the mother and *50%* from the father. 
+            These genes are organized into pairs of alleles, one from each parent. 
+            
+            A gene can be dominant, meaning it is expressed and determines a visible trait, or recessive, remaining unexpressed but still present in the genome. 
+            A recessive allele may manifest in future generations if both parents pass it on.
+
+            Five distinct genetic populations will be used in this simulation: **Americas, African, European, East Asian** and **South Asian**.
+
+            If you're curious to learn more about this simulation algorithm, you can refer to the code [here](https://github.com/Soptq/encDNA).
+            """
+        )
+
+        simulate_btn = gr.Button("Generate a random genetic allele", elem_id="input-box")
+
+        ethnicity_simulation_img = gr.Image(
+            visible=False,
+            container=False,
+            show_download_button=False,
+            mirror_webcam=False,
+            elem_id="img",
+        )
+
+        gr.Markdown(
+            """
+            The percentages you see above represent the two predominant genes transmitted to a progeny. Whether they are expressed or not is beyond the scope of this demo. 
+            Below is the vector representing the allele, generated above. 
+            """
+        )
+
+        clear_input_box = gr.Textbox(label="Unencrypted Allele:", visible=True)
+
+        gr.Markdown(
+            """
+            In practice, each individual is represented as a binary vector of size **1,059,079**, 
+            where each value corresponds to a genetic variation at a specific position on the chromosome: 1 represents the presence of a particular SNP, while 0 indicates its absence.
+            """
+        )
+
+        simulate_text = gr.Markdown(visible=False)
+
+        gr.Markdown("# Step 2: Encrypt the DNA on the Client Side")
+
+        gr.Markdown(
+            """
+            ⚠️ Important note: Encrypting an allele with more than a million values in FHE may take some time.
+            """
+        )
+
+        gen_key_btn = gr.Button("Generate the secret and public keys")
+        user_id_btn = gr.Number(value=USER_ID, visible=False)
+
+        # gr.HTML("<div style='height: 24px;'></div>")
+        encrypt_btn = gr.Button("Encrypt the data using the secret key")
+        encrypt_input_box = gr.Textbox(label="Encrypted Allele:", max_lines=15, visible=True)
+        send_btn = gr.Button("Send data to the server")
+
+        gr.Markdown("<hr />")
+        gr.Markdown("# Step 3: FHE Computation on the Server Side ")
+        gr.Markdown(
+            """
+            ⚠️ Important note: Processing such a large input in FHE may take some time, potentially up to 5 minutes for one allele of size *1,059,079*. 
+            To learn more about the selected ML model, check out this detailed [blog post](https://www.zama.ai/post/build-an-end-to-end-encrypted-23andme-genetic-testing-application-using-concrete-ml-fully-homomorphic-encryption).
+            """
+        )
+
+        run_fhe_btn = gr.Button("Run FHE on the server")
+        get_output_btn = gr.Button("Send data to the client")
+
+        gr.Markdown("<hr />")
+        gr.Markdown("# Step 4: Decrypt the Result on the Client Side")
+
+        decrypt_btn = gr.Button("Decrypt the data using the secret key")
+
+        pie_plot_output = gr.Image(
+            visible=False,
+            height=600,
+            width=600,
+            container=False,
+            show_download_button=False,
+            mirror_webcam=False,
+            elem_id="ie_plot_output",
+        )
+
+        gr.Markdown(
+            """
+            With FHE, the entire process remains encrypted end-to-end. Hence, you do not have to worry about data misuse, unauthorized analysis on your DNA, or data leaks. 
+            Using FHE, privacy is guaranteed, and trusting the server is no longer a concern.
+            """,
+            visible=False,
+        )
+
+        gr.Markdown("<hr />")
+        gr.Markdown(
+            """
+            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. 
+            Try it yourself and don't forget to star on [Github](https://github.com/zama-ai/concrete-ml) ⭐.
+            """
+        )
+        gr.Markdown()
+        gr.Markdown(
+            """
+            **Note that:** This space and the results produced by this simulation are for educational and illustrative purposes only. 
+            They are not intended to provide actual genetic analysis or be used as a substitute for a professional genetic testing.
+            """
+        )
+
+        ############################################################################# Click buttons
+
+        simulate_btn.click(
+            fn=simulate_allele_fn,
+            outputs=[clear_input_box, ethnicity_simulation_img, simulate_btn, simulate_text],
+        )
+
+        gen_key_btn.click(
+            key_gen_fn,
+            inputs=[user_id_btn],
+            outputs=[gen_key_btn],
+        )
+
+        encrypt_btn.click(
+            fn=encrypt_fn,
+            inputs=[user_id_btn],
+            outputs=[
+                encrypt_btn,
+                encrypt_input_box,
+                simulate_btn,
+                send_btn,
+                run_fhe_btn,
+                get_output_btn,
+                decrypt_btn,
+            ],
+        )
+
+        send_btn.click(
+            fn=send_input_fn,
+            inputs=[user_id_btn],
+            outputs=[send_btn, run_fhe_btn, get_output_btn, decrypt_btn],
+        )
+
+        run_fhe_btn.click(
+            fn=run_fhe_fn, inputs=[user_id_btn], outputs=[run_fhe_btn, get_output_btn, decrypt_btn]
+        )
+
+        get_output_btn.click(
+            fn=get_output_fn, inputs=[user_id_btn], outputs=[get_output_btn, decrypt_btn]
+        )
+
+        decrypt_btn.click(
+            fn=decrypt_fn, inputs=[user_id_btn], outputs=[decrypt_btn, pie_plot_output, user_id_btn]
+        )
+
+        demo.queue()
+        demo.launch()

server.py

@@ -0,0 +1,335 @@
+import time
+from glob import glob
+from pathlib import Path
+from typing import List
+
+from fastapi import FastAPI, File, Form, UploadFile
+from fastapi.responses import JSONResponse, Response
+from tqdm import tqdm
+from utils import *
+
+from concrete.ml.deployment import FHEModelClient, FHEModelServer
+
+# Load the FHE server
+
+# Initialize an instance of FastAPI
+app = FastAPI()
+
+
+# Define the default route
+@app.get("/")
+def root():
+    """
+    Root endpoint of the health prediction API.
+    Returns:
+        dict: The welcome message.
+    """
+    return {"message": "Welcome to your encrypted DNA testing use-case with FHE!"}
+
+
+@app.post("/send_input")
+def send_input(
+    user_id: str = Form(...), root_dir: str = Form(...), files: List[UploadFile] = File(...)
+):
+    """Send the inputs to the server."""
+
+    print("------------ Step 3.2: Send the data to the server")
+
+    print(f"{user_id=}, {root_dir=}, {len(files)=}")
+
+    SERVER_DIR = Path(root_dir) / f"{user_id}/server"
+
+    SERVER_KEY_SMOOTHER_MODULE_DIR = SERVER_DIR / KEY_SMOOTHER_MODULE_DIR
+    SERVER_KEY_BASE_MODULE_DIR = SERVER_DIR / KEY_BASE_MODULE_DIR
+    SERVER_ENCRYPTED_INPUT_DIR = SERVER_DIR / ENCRYPTED_INPUT_DIR
+
+    # Save the files using the above paths
+    with (SERVER_KEY_BASE_MODULE_DIR / "eval_key").open("wb") as eval_key_1:
+        eval_key_1.write(files[0].file.read())
+
+    with (SERVER_KEY_SMOOTHER_MODULE_DIR / "eval_key").open("wb") as eval_key_2:
+        eval_key_2.write(files[1].file.read())
+
+    print(f"{len(files)=}")
+    for i in tqdm(range(2, len(files))):
+        with (SERVER_ENCRYPTED_INPUT_DIR / f"encrypted_window_{i}").open("wb") as eval_key_2:
+            eval_key_2.write(files[i].file.read())
+
+
+@app.post("/run_fhe")
+def run_fhe(
+    user_id: str = Form(),
+    root_dir: str = Form(...),
+):
+    """Inference in FHE."""
+    print("------------ Step 4.2: Run in FHE on the Server Side")
+    print(f"{user_id=}, {root_dir=}")
+
+    SERVER_DIR = Path(root_dir) / f"{user_id}/server"
+
+    SERVER_KEY_SMOOTHER_MODULE_DIR = SERVER_DIR / KEY_SMOOTHER_MODULE_DIR
+    SERVER_KEY_BASE_MODULE_DIR = SERVER_DIR / KEY_BASE_MODULE_DIR
+    SERVER_ENCRYPTED_INPUT_DIR = SERVER_DIR / ENCRYPTED_INPUT_DIR
+    SERVER_ENCRYPTED_OUTPUT_DIR = SERVER_DIR / ENCRYPTED_OUTPUT_DIR
+
+    with (SERVER_KEY_BASE_MODULE_DIR / "eval_key").open("rb") as eval_key_1:
+        eval_key_base_module = eval_key_1.read()
+    assert isinstance(eval_key_base_module, bytes)
+
+    with (SERVER_KEY_SMOOTHER_MODULE_DIR / "eval_key").open("rb") as eval_key_2:
+        eval_key_smoother_module = eval_key_2.read()
+    assert isinstance(eval_key_smoother_module, bytes)
+
+    shared_base_modules_path = glob(f"{SHARED_BASE_MODULE_DIR}/model_*")
+    shared_base_modules_path = sorted(shared_base_modules_path, key=extract_model_number)
+    print(f"{len(shared_base_modules_path)=}")
+    assert len(shared_base_modules_path) == META["NW"]
+
+    client_encrypted_input_path = glob(f"{SERVER_ENCRYPTED_INPUT_DIR}/encrypted_window_*")
+    client_encrypted_input_path = sorted(client_encrypted_input_path, key=extract_model_number)
+    print(f"{len(client_encrypted_input_path)=}")
+    assert len(shared_base_modules_path) == META["NW"]
+
+    nb_total_iterations = META["NW"] * 2
+
+    start_time = time.time()
+    y_proba = []
+    for i, (model_path, encrypted_window_path) in tqdm(
+        enumerate(zip(shared_base_modules_path, client_encrypted_input_path))
+    ):
+        server = FHEModelServer(model_path)
+        with open(encrypted_window_path, "rb") as f:
+            encrypted_window = f.read()
+
+        encrypted_output = server.run(
+            encrypted_window, serialized_evaluation_keys=eval_key_base_module
+        )
+        assert isinstance(encrypted_output, bytes)
+
+        client = FHEModelClient(model_path, key_dir=model_path)
+        decrypted_output = client.deserialize_decrypt_dequantize(encrypted_output)
+
+        with (SERVER_ENCRYPTED_OUTPUT_DIR / f"decrypted_window_{i}").open("wb") as f:
+            f.write(encrypted_window)
+
+        y_proba.append(decrypted_output)
+
+        with open(FHE_COMPUTATION_TIMELINE, "w", encoding="utf-8") as f:
+            f.write(f"{time.time() - start_time:.2f} seconds ({(i + 1)/nb_total_iterations:.0%})")
+
+    client = FHEModelClient(SHARED_SMOOTHER_MODULE_DIR, key_dir=SHARED_SMOOTHER_MODULE_DIR)
+    server = FHEModelServer(SHARED_SMOOTHER_MODULE_DIR)
+
+    y_proba = numpy.transpose(numpy.array(y_proba), (1, 0, 2))
+    y_proba = y_proba.astype(numpy.int8)
+    print(f"{y_proba.shape=}, {type(y_proba)}")
+
+    X_slide, _ = slide_window(y_proba, META["SS"])
+
+    yhat_encrypted = []
+    for i in tqdm(range(len(X_slide))):
+        input = X_slide[i].reshape(1, -1)
+        encrypted_input = client.quantize_encrypt_serialize(input)
+        encrypted_output = server.run(
+            encrypted_input, serialized_evaluation_keys=eval_key_smoother_module
+        )
+        # output = client.deserialize_decrypt_dequantize(encrypted_output)
+        # y_pred = numpy.argmax(output, axis=-1)[0]
+        yhat_encrypted.append(encrypted_output)
+
+        with open(FHE_COMPUTATION_TIMELINE, "w", encoding="utf-8") as f:
+            f.write(f"{time.time() - start_time:.2f} seconds ({(i + 1)/nb_total_iterations:.0%})")
+
+    write_pickle(SERVER_ENCRYPTED_OUTPUT_DIR / "encrypted_final_output.pkl", yhat_encrypted)
+
+    fhe_execution_time = round(time.time() - start_time, 2)
+
+    return JSONResponse(content=fhe_execution_time)
+
+
+@app.post("/run_fhe_stage1")
+def run_fhe_stage1(
+    user_id: str = Form(),
+    root_dir: str = Form(...),
+):
+    """Inference in FHE."""
+    print("------------ Step 4.2: Run in FHE on the Server Side")
+    print(f"{user_id=}, {root_dir=}")
+
+    SERVER_DIR = Path(root_dir) / f"{user_id}/server"
+
+    SERVER_KEY_SMOOTHER_MODULE_DIR = SERVER_DIR / KEY_SMOOTHER_MODULE_DIR
+    SERVER_KEY_BASE_MODULE_DIR = SERVER_DIR / KEY_BASE_MODULE_DIR
+    SERVER_ENCRYPTED_INPUT_DIR = SERVER_DIR / ENCRYPTED_INPUT_DIR
+    SERVER_ENCRYPTED_OUTPUT_DIR = SERVER_DIR / ENCRYPTED_OUTPUT_DIR
+
+    with (SERVER_KEY_BASE_MODULE_DIR / "eval_key").open("rb") as eval_key_1:
+        eval_key_base_module = eval_key_1.read()
+    assert isinstance(eval_key_base_module, bytes)
+
+    with (SERVER_KEY_SMOOTHER_MODULE_DIR / "eval_key").open("rb") as eval_key_2:
+        eval_key_smoother_module = eval_key_2.read()
+    assert isinstance(eval_key_smoother_module, bytes)
+
+    shared_base_modules_path = glob(f"{SHARED_BASE_MODULE_DIR}/model_*")
+    shared_base_modules_path = sorted(shared_base_modules_path, key=extract_model_number)
+    print(f"{len(shared_base_modules_path)=}")
+    assert len(shared_base_modules_path) == META["NW"]
+
+    client_encrypted_input_path = glob(f"{SERVER_ENCRYPTED_INPUT_DIR}/encrypted_window_*")
+    client_encrypted_input_path = sorted(client_encrypted_input_path, key=extract_model_number)
+    print(f"{len(client_encrypted_input_path)=}")
+    assert len(shared_base_modules_path) == META["NW"]
+
+    start = time.time()
+    y_proba = []
+    for i, (model_path, encrypted_window_path) in tqdm(
+        enumerate(zip(shared_base_modules_path, client_encrypted_input_path))
+    ):
+        server = FHEModelServer(model_path)
+        with open(encrypted_window_path, "rb") as f:
+            encrypted_window = f.read()
+
+        encrypted_output = server.run(
+            encrypted_window, serialized_evaluation_keys=eval_key_base_module
+        )
+        assert isinstance(encrypted_output, bytes)
+
+        client = FHEModelClient(model_path, key_dir=model_path)
+        decrypted_output = client.deserialize_decrypt_dequantize(encrypted_output)
+
+        with (SERVER_ENCRYPTED_OUTPUT_DIR / f"decrypted_window_{i}").open("wb") as f:
+            f.write(encrypted_window)
+
+        y_proba.append(decrypted_output)
+
+    client = FHEModelClient(SHARED_SMOOTHER_MODULE_DIR, key_dir=SHARED_SMOOTHER_MODULE_DIR)
+    server = FHEModelServer(SHARED_SMOOTHER_MODULE_DIR)
+
+    y_proba = numpy.transpose(numpy.array(y_proba), (1, 0, 2))
+    y_proba = y_proba.astype(numpy.int8)
+    print(f"{y_proba.shape=}, {type(y_proba)}")
+
+    X_slide, _ = slide_window(y_proba, META["SS"])
+
+    yhat_encrypted = []
+    for i in tqdm(range(len(X_slide))):
+        input = X_slide[i].reshape(1, -1)
+        encrypted_input = client.quantize_encrypt_serialize(input)
+        encrypted_output = server.run(
+            encrypted_input, serialized_evaluation_keys=eval_key_smoother_module
+        )
+        # output = client.deserialize_decrypt_dequantize(encrypted_output)
+        # y_pred = numpy.argmax(output, axis=-1)[0]
+        yhat_encrypted.append(encrypted_output)
+
+    write_pickle(SERVER_ENCRYPTED_OUTPUT_DIR / "encrypted_final_output.pkl", yhat_encrypted)
+
+    fhe_execution_time = round(time.time() - start, 2)
+
+    return JSONResponse(content=fhe_execution_time)
+
+
+@app.post("/run_fhe_stage2")
+def run_fhe_stage2(
+    user_id: str = Form(),
+    root_dir: str = Form(...),
+):
+    """Inference in FHE."""
+    print("------------ Step 4.2: Run in FHE on the Server Side")
+    print(f"{user_id=}, {root_dir=}")
+
+    SERVER_DIR = Path(root_dir) / f"{user_id}/server"
+
+    SERVER_KEY_SMOOTHER_MODULE_DIR = SERVER_DIR / KEY_SMOOTHER_MODULE_DIR
+    SERVER_KEY_BASE_MODULE_DIR = SERVER_DIR / KEY_BASE_MODULE_DIR
+    SERVER_ENCRYPTED_INPUT_DIR = SERVER_DIR / ENCRYPTED_INPUT_DIR
+    SERVER_ENCRYPTED_OUTPUT_DIR = SERVER_DIR / ENCRYPTED_OUTPUT_DIR
+
+    with (SERVER_KEY_BASE_MODULE_DIR / "eval_key").open("rb") as eval_key_1:
+        eval_key_base_module = eval_key_1.read()
+    assert isinstance(eval_key_base_module, bytes)
+
+    with (SERVER_KEY_SMOOTHER_MODULE_DIR / "eval_key").open("rb") as eval_key_2:
+        eval_key_smoother_module = eval_key_2.read()
+    assert isinstance(eval_key_smoother_module, bytes)
+
+    shared_base_modules_path = glob(f"{SHARED_BASE_MODULE_DIR}/model_*")
+    shared_base_modules_path = sorted(shared_base_modules_path, key=extract_model_number)
+    print(f"{len(shared_base_modules_path)=}")
+    assert len(shared_base_modules_path) == META["NW"]
+
+    client_encrypted_input_path = glob(f"{SERVER_ENCRYPTED_INPUT_DIR}/encrypted_window_*")
+    client_encrypted_input_path = sorted(client_encrypted_input_path, key=extract_model_number)
+    print(f"{len(client_encrypted_input_path)=}")
+    assert len(shared_base_modules_path) == META["NW"]
+
+    start = time.time()
+    y_proba = []
+    for i, (model_path, encrypted_window_path) in tqdm(
+        enumerate(zip(shared_base_modules_path, client_encrypted_input_path))
+    ):
+        server = FHEModelServer(model_path)
+        with open(encrypted_window_path, "rb") as f:
+            encrypted_window = f.read()
+
+        encrypted_output = server.run(
+            encrypted_window, serialized_evaluation_keys=eval_key_base_module
+        )
+        assert isinstance(encrypted_output, bytes)
+
+        client = FHEModelClient(model_path, key_dir=model_path)
+        decrypted_output = client.deserialize_decrypt_dequantize(encrypted_output)
+
+        with (SERVER_ENCRYPTED_OUTPUT_DIR / f"decrypted_window_{i}").open("wb") as f:
+            f.write(encrypted_window)
+
+        y_proba.append(decrypted_output)
+
+    client = FHEModelClient(SHARED_SMOOTHER_MODULE_DIR, key_dir=SHARED_SMOOTHER_MODULE_DIR)
+    server = FHEModelServer(SHARED_SMOOTHER_MODULE_DIR)
+
+    y_proba = numpy.transpose(numpy.array(y_proba), (1, 0, 2))
+    y_proba = y_proba.astype(numpy.int8)
+    print(f"{y_proba.shape=}, {type(y_proba)}")
+
+    X_slide, _ = slide_window(y_proba, META["SS"])
+
+    yhat_encrypted = []
+    for i in tqdm(range(len(X_slide))):
+        input = X_slide[i].reshape(1, -1)
+        encrypted_input = client.quantize_encrypt_serialize(input)
+        encrypted_output = server.run(
+            encrypted_input, serialized_evaluation_keys=eval_key_smoother_module
+        )
+        # output = client.deserialize_decrypt_dequantize(encrypted_output)
+        # y_pred = numpy.argmax(output, axis=-1)[0]
+        yhat_encrypted.append(encrypted_output)
+
+    write_pickle(SERVER_ENCRYPTED_OUTPUT_DIR / "encrypted_final_output.pkl", yhat_encrypted)
+
+    fhe_execution_time = round(time.time() - start, 2)
+
+    return JSONResponse(content=fhe_execution_time)
+
+
+@app.post("/get_output")
+def get_output(user_id: str = Form(), root_dir: str = Form()):
+    """Retrieve the encrypted output from the server."""
+
+    print("\nStep 5.2: Get the output from the server ............\n")
+    SERVER_DIR = Path(root_dir) / f"{user_id}/server"
+    SERVER_ENCRYPTED_OUTPUT_DIR = SERVER_DIR / ENCRYPTED_OUTPUT_DIR
+
+    yhat_encrypted = load_pickle(SERVER_ENCRYPTED_OUTPUT_DIR / "encrypted_final_output.pkl")
+
+    CLIENT_DIR = Path(root_dir) / f"{user_id}/client"
+    CLIENT_ENCRYPTED_OUTPUT_DIR = CLIENT_DIR / ENCRYPTED_OUTPUT_DIR
+
+    write_pickle(CLIENT_ENCRYPTED_OUTPUT_DIR / "encrypted_final_output.pkl", yhat_encrypted)
+    assert len(yhat_encrypted) == META["NW"]
+
+    time.sleep(1)
+
+    # Send the encrypted output
+    return Response("OK")

utils.py

@@ -0,0 +1,179 @@
+import pickle
+import random
+import shutil
+from collections import Counter
+from pathlib import Path
+
+import numpy
+
+import zipfile
+
+
+SERVER_URL = "http://localhost:8000/"
+
+INPUT_BROWSER_LIMIT = 550
+
+DATA_DIR = Path("./data")
+DEPLOYMENT_DIR = Path("./deployment")
+ROOT_DIR = DEPLOYMENT_DIR / "users"
+
+SHARED_BASE_MODULE_DIR = DEPLOYMENT_DIR / "base_modules"
+SHARED_SMOOTHER_MODULE_DIR = DEPLOYMENT_DIR / "smoother_module"
+
+KEY_SMOOTHER_MODULE_DIR = "EvaluationKey_Smoother"
+KEY_BASE_MODULE_DIR = "EvaluationKey_Base_Modules"
+ENCRYPTED_INPUT_DIR = "Encrypt_Input"
+ENCRYPTED_OUTPUT_DIR = "Encrypt_Output"
+FHE_COMPUTATION_TIMELINE = Path("server_fhe_computation_timeline.txt")
+
+LABELS = ["European", "African", "Americas", "East Asian", "South Asian"]
+
+ID_POPULATION = {0: "European", 3: "African", 2: "Americas", 1: "East Asian", 4: "South Asian"}
+
+POPULATION_ID = {"European": 0, "African": 3, "Americas": 2, "East Asian": 1, "South Asian": 4}
+
+COLORS = ["#FFD208", "#FFE46C", "#FFED9C", "#FFF6CE", "#FFD9A0"]
+
+# load_pickle("data/meta_dict.pkl")
+META = {"A": 5, "C": 1059079, "M": 10589, "NW": 100, "CT": 1059, "CTR": 0.1, "WSCM": 0.2, "SS": 75}
+
+BUILD_GENS = [1, 2, 4, 6, 8, 12, 16, 24, 32, 48]
+
+
+def load_pickle_from_zip(zip_path, file_name):
+    """
+    Load a pickle file from within a zip archive.
+    """
+    with zipfile.ZipFile(zip_path, 'r') as z:
+        with z.open(file_name) as f:
+            return pickle.load(f)
+            
+
+def generate_weighted_percentages():
+    dominant_percentage = random.randint(50, 70)
+    remaining_percentage = 100 - dominant_percentage
+    other_percentages = [random.random() for _ in range(4)]
+
+    total = sum(other_percentages)
+    other_percentages = [round(p / total * remaining_percentage, 2) for p in other_percentages]
+
+    percentages = [dominant_percentage] + other_percentages
+
+    # Adjust the total to be exactly 100 (if rounding errors occurred)
+    diff = round(100 - sum(percentages), 2)
+    if diff != 0:
+        percentages[0] += diff  # Adjust the dominant percentage to make the total 100
+
+    return percentages
+
+
+def select_random_ancestors():
+    ancestors = list(ID_POPULATION.keys())
+    random.shuffle(ancestors)
+    return ancestors
+
+
+def read_pickle(path):
+    with open(path, "rb") as f:
+        data = pickle.load(f)
+    return data
+
+
+def compute_distribution(y, size=5):
+    y_pred = numpy.zeros(size)
+    for k, v in Counter(y).items():
+        y_pred[k] = v / len(y)
+    return y_pred
+
+
+def slide_window(data, smooth_win_size, y=None):
+    N, W, A = data.shape
+
+    pad = (smooth_win_size + 1) // 2
+    data_padded = numpy.pad(data, ((0, 0), (pad, pad), (0, 0)), mode="reflect")
+    X_slide = numpy.lib.stride_tricks.sliding_window_view(data_padded, (1, smooth_win_size, A))
+    X_slide = X_slide[:, :W, :].reshape(N * W, -1)
+    y_slide = None if y is None else y.reshape(N * W)
+
+    return X_slide, y_slide
+
+
+# def read_vcf(vcf_file):
+#     return allel.read_vcf(vcf_file, region=None, fields="*")
+
+
+def clean_dir(directory):
+    """Remove the specified directory if it exists."""
+    if directory.exists() and directory.is_dir():
+        print(f"Removing existing model directory: {directory}")
+        shutil.rmtree(directory)
+
+
+def process_data_for_base_modules(meta, X_t):
+
+    n_windows = meta["NW"]  # meta["C"] // meta["M"]
+    context = meta["CT"]  # int(meta["M"] * meta['CTR'])
+
+    if context != 0.0:
+        pad_left = numpy.flip(X_t[:, 0:context], axis=1)
+        pad_right = numpy.flip(X_t[:, -context:], axis=1)
+        X_t = numpy.concatenate([pad_left, X_t, pad_right], axis=1)
+
+    M_ = meta["M"] + 2 * context
+    idx = numpy.arange(0, meta["C"], meta["M"])[:-2]
+    X_b = numpy.lib.stride_tricks.sliding_window_view(X_t, M_, axis=1)[:, idx, :]
+    rem = meta["C"] - meta["M"] * n_windows
+
+    # print(f"{X_t.shape=} -> {X_b.shape=} | {n_windows=}, {context=}, {M_=}, {rem=}")
+
+    return X_b, n_windows, M_, rem
+
+
+def extract_model_number(path):
+    try:
+        return int(path.split("_")[-1])
+    except (ValueError, IndexError):
+        print(f"Error: Unable to extract model number from path: {path}")
+        return None
+
+
+def is_none(obj) -> bool:
+    """
+    Check if the object is None.
+    Args:
+        obj (any): The input to be checked.
+    Returns:
+        bool: True if the object is None or empty, False otherwise.
+    """
+    return obj is None or (obj is not None and (hasattr(obj, "__len__") and len(obj) == 0))
+
+
+def load_pickle(path: str) -> numpy.array:
+    """Load data.
+
+    Args:
+        path (str):
+
+    Returns:
+        Dict: The genome.
+    """
+    with open(path, "rb") as f:
+        data = pickle.load(f)
+    return data
+
+
+def write_pickle(path: str, data) -> numpy.array:
+    with open(path, "wb") as f:
+        pickle.dump(data, f)
+
+
+def write_bytes(path, data):
+    """Save binary data."""
+    with path.open("wb") as f:
+        f.write(data)
+
+
+def read_bytes(path):
+    """Load data from a binary file."""
+    with path.open("rb") as f:
+        return f.read()