|
"""Backend functions used in the app.""" |
|
|
|
import os |
|
import shutil |
|
import gradio as gr |
|
import numpy |
|
import requests |
|
import pickle |
|
import pandas |
|
from itertools import chain |
|
|
|
from settings import ( |
|
SERVER_URL, |
|
FHE_KEYS, |
|
CLIENT_FILES, |
|
SERVER_FILES, |
|
DEPLOYMENT_PATH, |
|
INITIAL_INPUT_SHAPE, |
|
INPUT_INDEXES, |
|
INPUT_SLICES, |
|
PRE_PROCESSOR_USER_PATH, |
|
PRE_PROCESSOR_THIRD_PARTY_PATH, |
|
CLIENT_TYPES, |
|
) |
|
|
|
from utils.client_server_interface import MultiInputsFHEModelClient |
|
|
|
|
|
with PRE_PROCESSOR_USER_PATH.open('rb') as file: |
|
PRE_PROCESSOR_USER = pickle.load(file) |
|
|
|
with PRE_PROCESSOR_THIRD_PARTY_PATH.open('rb') as file: |
|
PRE_PROCESSOR_THIRD_PARTY = pickle.load(file) |
|
|
|
|
|
def shorten_bytes_object(bytes_object, limit=500): |
|
"""Shorten the input bytes object to a given length. |
|
|
|
Encrypted data is too large for displaying it in the browser using Gradio. This function |
|
provides a shorten representation of it. |
|
|
|
Args: |
|
bytes_object (bytes): The input to shorten |
|
limit (int): The length to consider. Default to 500. |
|
|
|
Returns: |
|
str: Hexadecimal string shorten representation of the input byte object. |
|
|
|
""" |
|
|
|
shift = 100 |
|
return bytes_object[shift : limit + shift].hex() |
|
|
|
|
|
def clean_temporary_files(n_keys=20): |
|
"""Clean keys and encrypted inputs. |
|
|
|
A maximum of n_keys keys and associated temporary files are allowed to be stored. Once this |
|
limit is reached, the oldest files are deleted. |
|
|
|
Args: |
|
n_keys (int): The maximum number of keys and associated files to be stored. Default to 20. |
|
|
|
""" |
|
|
|
key_dirs = sorted(FHE_KEYS.iterdir(), key=os.path.getmtime) |
|
|
|
|
|
user_ids = [] |
|
if len(key_dirs) > n_keys: |
|
n_keys_to_delete = len(key_dirs) - n_keys |
|
for key_dir in key_dirs[:n_keys_to_delete]: |
|
user_ids.append(key_dir.name) |
|
shutil.rmtree(key_dir) |
|
|
|
|
|
client_files = CLIENT_FILES.iterdir() |
|
server_files = SERVER_FILES.iterdir() |
|
|
|
|
|
for file in chain(client_files, server_files): |
|
for user_id in user_ids: |
|
if user_id in file.name: |
|
file.unlink() |
|
|
|
|
|
def _get_client(client_id, client_type): |
|
"""Get the client API. |
|
|
|
Args: |
|
client_id (int): The client ID to consider. |
|
client_type (str): The type of user to consider (either 'user', 'bank' or 'third_party'). |
|
|
|
Returns: |
|
FHEModelClient: The client API. |
|
""" |
|
key_dir = FHE_KEYS / f"{client_type}_{client_id}" |
|
|
|
return MultiInputsFHEModelClient(DEPLOYMENT_PATH, key_dir=key_dir, nb_inputs=len(CLIENT_TYPES)) |
|
|
|
|
|
def _keygen(client_id, client_type): |
|
"""Generate the private key associated to a client. |
|
|
|
Args: |
|
client_id (int): The client ID to consider. |
|
client_type (str): The type of client to consider (either 'user', 'bank' or 'third_party'). |
|
""" |
|
|
|
clean_temporary_files() |
|
|
|
|
|
client = _get_client(client_id, client_type) |
|
|
|
|
|
client.generate_private_and_evaluation_keys(force=True) |
|
|
|
|
|
|
|
|
|
evaluation_key = client.get_serialized_evaluation_keys() |
|
|
|
|
|
|
|
evaluation_key_path = _get_client_file_path("evaluation_key", client_id, client_type) |
|
|
|
with evaluation_key_path.open("wb") as evaluation_key_file: |
|
evaluation_key_file.write(evaluation_key) |
|
|
|
|
|
def _send_input(client_id, client_type): |
|
"""Send the encrypted inputs as well as the evaluation key to the server. |
|
|
|
Args: |
|
client_id (int): The client ID to consider. |
|
client_type (str): The type of client to consider (either 'user', 'bank' or 'third_party'). |
|
""" |
|
|
|
evaluation_key_path = _get_client_file_path("evaluation_key", client_id, client_type) |
|
encrypted_input_path = _get_client_file_path("encrypted_inputs", client_id, client_type) |
|
|
|
|
|
data = { |
|
"client_id": client_id, |
|
"client_type": client_type, |
|
} |
|
|
|
files = [ |
|
("files", open(encrypted_input_path, "rb")), |
|
("files", open(evaluation_key_path, "rb")), |
|
] |
|
|
|
|
|
url = SERVER_URL + "send_input" |
|
with requests.post( |
|
url=url, |
|
data=data, |
|
files=files, |
|
) as response: |
|
return response.ok |
|
|
|
|
|
def _get_client_file_path(name, client_id, client_type): |
|
"""Get the correct temporary file path for the client. |
|
|
|
Args: |
|
name (str): The desired file name (either 'evaluation_key' or 'encrypted_inputs'). |
|
client_id (int): The client ID to consider. |
|
client_type (str): The type of user to consider (either 'user', 'bank' or 'third_party'). |
|
|
|
Returns: |
|
pathlib.Path: The file path. |
|
""" |
|
return CLIENT_FILES / f"{name}_{client_type}_{client_id}" |
|
|
|
|
|
def _keygen_encrypt_send(inputs, client_type): |
|
"""Encrypt the given inputs for a specific client. |
|
|
|
Args: |
|
inputs (numpy.ndarray): The inputs to encrypt. |
|
client_type (str): The type of client to consider (either 'user', 'bank' or 'third_party'). |
|
|
|
Returns: |
|
client_id, encrypted_inputs_short (int, bytes): Integer ID representing the current client |
|
and a byte short representation of the encrypted input to send. |
|
""" |
|
|
|
client_id = numpy.random.randint(0, 2**32) |
|
|
|
_keygen(client_id, client_type) |
|
|
|
|
|
client = _get_client(client_id, client_type) |
|
|
|
|
|
|
|
|
|
encrypted_inputs = client.quantize_encrypt_serialize_multi_inputs( |
|
inputs, |
|
input_index=INPUT_INDEXES[client_type], |
|
initial_input_shape=INITIAL_INPUT_SHAPE, |
|
input_slice=INPUT_SLICES[client_type], |
|
) |
|
|
|
|
|
|
|
encrypted_inputs_path = _get_client_file_path("encrypted_inputs", client_id, client_type) |
|
|
|
with encrypted_inputs_path.open("wb") as encrypted_inputs_file: |
|
encrypted_inputs_file.write(encrypted_inputs) |
|
|
|
|
|
encrypted_inputs_short = shorten_bytes_object(encrypted_inputs) |
|
|
|
_send_input(client_id, client_type) |
|
|
|
|
|
return client_id, encrypted_inputs_short |
|
|
|
|
|
def pre_process_keygen_encrypt_send_user(*inputs): |
|
"""Pre-process the given inputs for a specific client. |
|
|
|
Args: |
|
*inputs (Tuple[numpy.ndarray]): The inputs to pre-process. |
|
|
|
Returns: |
|
(int, bytes): Integer ID representing the current client and a byte short representation of |
|
the encrypted input to send. |
|
""" |
|
gender, bool_inputs, num_children, num_family, total_income, age, income_type, education_type, \ |
|
family_status, occupation_type, housing_type = inputs |
|
|
|
|
|
|
|
gender = gender == "Male" |
|
|
|
|
|
own_car = "Car" in bool_inputs |
|
own_property = "Property" in bool_inputs |
|
work_phone = "Work phone" in bool_inputs |
|
phone = "Phone" in bool_inputs |
|
email = "Email" in bool_inputs |
|
|
|
user_inputs = pandas.DataFrame({ |
|
"Gender": [gender], |
|
"Own_car": [own_car], |
|
"Own_property": [own_property], |
|
"Work_phone": [work_phone], |
|
"Phone": [phone], |
|
"Email": [email], |
|
"Num_children": num_children, |
|
"Num_family": num_family, |
|
"Total_income": total_income, |
|
"Age": age, |
|
"Income_type": income_type, |
|
"Education_type": education_type, |
|
"Family_status": family_status, |
|
"Occupation_type": occupation_type, |
|
"Housing_type": housing_type, |
|
}) |
|
|
|
preprocessed_user_inputs = PRE_PROCESSOR_USER.transform(user_inputs) |
|
|
|
return _keygen_encrypt_send(preprocessed_user_inputs, "user") |
|
|
|
|
|
def pre_process_keygen_encrypt_send_bank(*inputs): |
|
"""Pre-process the given inputs for a specific client. |
|
|
|
Args: |
|
*inputs (Tuple[numpy.ndarray]): The inputs to pre-process. |
|
|
|
Returns: |
|
(int, bytes): Integer ID representing the current client and a byte short representation of |
|
the encrypted input to send. |
|
""" |
|
account_length = inputs[0] |
|
|
|
return _keygen_encrypt_send(account_length, "bank") |
|
|
|
|
|
def pre_process_keygen_encrypt_send_third_party(*inputs): |
|
"""Pre-process the given inputs for a specific client. |
|
|
|
Args: |
|
*inputs (Tuple[numpy.ndarray]): The inputs to pre-process. |
|
|
|
Returns: |
|
(int, bytes): Integer ID representing the current client and a byte short representation of |
|
the encrypted input to send. |
|
""" |
|
employed, years_employed = inputs |
|
|
|
|
|
unemployed = employed == "No" |
|
|
|
third_party_inputs = pandas.DataFrame({ |
|
"Unemployed": [unemployed], |
|
"Years_employed": [years_employed], |
|
}) |
|
|
|
preprocessed_third_party_inputs = PRE_PROCESSOR_THIRD_PARTY.transform(third_party_inputs) |
|
|
|
return _keygen_encrypt_send(preprocessed_third_party_inputs, "third_party") |
|
|
|
|
|
def run_fhe(user_id, bank_id, third_party_id): |
|
"""Run the model on the encrypted inputs previously sent using FHE. |
|
|
|
Args: |
|
user_id (int): The user ID to consider. |
|
bank_id (int): The bank ID to consider. |
|
third_party_id (int): The third party ID to consider. |
|
""" |
|
|
|
|
|
|
|
data = { |
|
"user_id": user_id, |
|
"bank_id": bank_id, |
|
"third_party_id": third_party_id, |
|
} |
|
|
|
|
|
url = SERVER_URL + "run_fhe" |
|
with requests.post( |
|
url=url, |
|
data=data, |
|
) as response: |
|
if response.ok: |
|
return response.json() |
|
else: |
|
raise gr.Error("Please wait for the inputs to be sent to the server.") |
|
|
|
|
|
def get_output(user_id, bank_id, third_party_id): |
|
"""Retrieve the encrypted output. |
|
|
|
Args: |
|
user_id (int): The user ID to consider. |
|
bank_id (int): The bank ID to consider. |
|
third_party_id (int): The third party ID to consider. |
|
""" |
|
data = { |
|
"user_id": user_id, |
|
"bank_id": bank_id, |
|
"third_party_id": third_party_id, |
|
} |
|
|
|
|
|
url = SERVER_URL + "get_output" |
|
with requests.post( |
|
url=url, |
|
data=data, |
|
) as response: |
|
if response.ok: |
|
encrypted_output = response.content |
|
|
|
|
|
|
|
|
|
encrypted_output_path = _get_client_file_path("encrypted_output", user_id + bank_id + third_party_id, "output") |
|
|
|
with encrypted_output_path.open("wb") as encrypted_output_file: |
|
encrypted_output_file.write(encrypted_output) |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
return None |
|
else: |
|
raise gr.Error("Please wait for the FHE execution to be completed.") |
|
|
|
|
|
def decrypt_output(user_id, bank_id, third_party_id): |
|
"""Decrypt the result. |
|
|
|
Args: |
|
user_id (int): The user ID to consider. |
|
bank_id (int): The bank ID to consider. |
|
third_party_id (int): The third party ID to consider. |
|
|
|
Returns: |
|
output(numpy.ndarray): The decrypted output |
|
|
|
""" |
|
|
|
encrypted_output_path = _get_client_file_path("encrypted_output", user_id + bank_id + third_party_id, "output") |
|
|
|
if not encrypted_output_path.is_file(): |
|
raise gr.Error("Please run the FHE execution first.") |
|
|
|
|
|
with encrypted_output_path.open("rb") as encrypted_output_file: |
|
encrypted_output_proba = encrypted_output_file.read() |
|
|
|
|
|
client = _get_client(user_id, "user") |
|
|
|
|
|
output_proba = client.deserialize_decrypt_dequantize(encrypted_output_proba) |
|
|
|
|
|
output = numpy.argmax(output_proba, axis=1) |
|
|
|
return output |