Spaces:
Sleeping
Sleeping
| import numpy as np | |
| import pandas as pd | |
| import requests | |
| import json | |
| import pickle | |
| import io | |
| import gradio as gr | |
| from fhe_utils import FHEClient | |
| from common import ( | |
| ROOT_DIR, | |
| KEYS_PATH, | |
| DUNEDIN_PACE_DATA_PATH, | |
| SERVER_URL, | |
| SCALE, | |
| get_tmp_file_path | |
| ) | |
| DUNEDIN_PACE_DATA = json.loads(open(DUNEDIN_PACE_DATA_PATH).read()) | |
| def get_client(user_id): | |
| return FHEClient( | |
| ROOT_DIR / "client.zip", | |
| scale=SCALE, | |
| dunedin_pace_data=DUNEDIN_PACE_DATA, | |
| key_dir=KEYS_PATH / f"{user_id}", | |
| ) | |
| def generate_fhe_key(): | |
| user_id = np.random.randint(0, 2**32) | |
| client = get_client(user_id) | |
| client.generate_private_and_evaluation_keys(force=True) | |
| evaluation_key = client.get_serialized_evaluation_keys() | |
| evaluation_key_path = get_tmp_file_path(user_id, "evaluation_key", "client") | |
| with evaluation_key_path.open("wb") as fw: | |
| fw.write(evaluation_key) | |
| return (user_id, True) | |
| def save_to_file(user_id, file_name, data): | |
| file_path = get_tmp_file_path(user_id, file_name) | |
| with file_path.open("wb") as fw: | |
| fw.write(data) | |
| def get_bytes(obj): | |
| obj_bytes = pickle.dumps(obj) | |
| return io.BytesIO(obj_bytes) | |
| def get_pace_estimation(client, evaluation_key_path, betas_sample): | |
| betas_sample_es, weights_int_es, intercept_int_es = client.encrypt_serialize(betas_sample) | |
| files = [ | |
| ("files", betas_sample_es), | |
| ("files", weights_int_es), | |
| ("files", intercept_int_es), | |
| ("files", open(evaluation_key_path, "rb")), | |
| ] | |
| resp = requests.post( | |
| f"{SERVER_URL}/pace_estimation", | |
| data={}, | |
| files=files | |
| ) | |
| pace_es = resp.content | |
| return client.deserialize_decrypt_post_process(pace_es) | |
| def get_pace_estimations(user_id, betas_file): | |
| evaluation_key_path = get_tmp_file_path(user_id, "evaluation_key") | |
| try: | |
| beta_values = pd.read_csv(betas_file, delimiter="\t", index_col=0) | |
| except Exception: | |
| return (None, gr.update(visible=False), "Non valid CSV") | |
| if len(beta_values.columns) == 0: | |
| return (None, gr.update(visible=False), "[DATA ERROR]: No samples provided") | |
| for base_feature in DUNEDIN_PACE_DATA['probes']: | |
| if base_feature not in beta_values.index: | |
| return (None, None, f"[DATA ERROR]: Missing base probe '{base_feature}'") | |
| probes_mean = { | |
| DUNEDIN_PACE_DATA['normalization_probes'][i]: DUNEDIN_PACE_DATA['normalization_means'][i] | |
| for i in range(len(DUNEDIN_PACE_DATA['normalization_probes'])) | |
| } | |
| missing_probes = [_ for _ in DUNEDIN_PACE_DATA['normalization_probes'] if _ not in beta_values.index] | |
| missing_probes_data = [[probes_mean[_]] * len(beta_values.columns) for _ in missing_probes] | |
| missing_betas = pd.DataFrame(missing_probes_data, index=missing_probes, columns=beta_values.columns) | |
| filtered_beta_values = pd.concat([beta_values, missing_betas]).loc[DUNEDIN_PACE_DATA['normalization_probes']] | |
| client = get_client(user_id) | |
| pace_data = [ | |
| (column_name, get_pace_estimation( | |
| client, evaluation_key_path, | |
| filtered_beta_values[column_name].values | |
| )) | |
| for column_name in filtered_beta_values.columns | |
| ] | |
| return (pace_data, gr.update(visible=True), "Success") |