Spaces:
Runtime error
Runtime error
"""Server that will listen for GET and POST requests from the client.""" | |
import time | |
from typing import List | |
from fastapi import FastAPI, File, Form, UploadFile | |
from fastapi.responses import JSONResponse, Response | |
from utils import DEPLOYMENT_DIR, SERVER_DIR # pylint: disable=no-name-in-module) | |
from concrete.ml.deployment import FHEModelServer | |
# Initialize an instance of FastAPI | |
app = FastAPI() | |
# Define the default route | |
def root(): | |
""" | |
Root endpoint of the health prediction API. | |
Returns: | |
dict: The welcome message. | |
""" | |
return {"message": "Welcome to your disease prediction with FHE!"} | |
def send_input( | |
user_id: str = Form(), | |
files: List[UploadFile] = File(), | |
): | |
"""Send the inputs to the server.""" | |
print("\nSend the data to the server ............\n") | |
# Receive the Client's files (Evaluation key + Encrypted symptoms) | |
evaluation_key_path = SERVER_DIR / f"{user_id}_valuation_key" | |
encrypted_input_path = SERVER_DIR / f"{user_id}_encrypted_input" | |
# Save the files using the above paths | |
with encrypted_input_path.open("wb") as encrypted_input, evaluation_key_path.open( | |
"wb" | |
) as evaluation_key: | |
encrypted_input.write(files[0].file.read()) | |
evaluation_key.write(files[1].file.read()) | |
def run_fhe( | |
user_id: str = Form(), | |
): | |
"""Inference in FHE.""" | |
print("\nRun in FHE in the server ............\n") | |
evaluation_key_path = SERVER_DIR / f"{user_id}_valuation_key" | |
encrypted_input_path = SERVER_DIR / f"{user_id}_encrypted_input" | |
# Read the files (Evaluation key + Encrypted symptoms) using the above paths | |
with encrypted_input_path.open("rb") as encrypted_output_file, evaluation_key_path.open( | |
"rb" | |
) as evaluation_key_file: | |
encrypted_output = encrypted_output_file.read() | |
evaluation_key = evaluation_key_file.read() | |
# Load the FHE server and the model | |
fhe_server = FHEModelServer(DEPLOYMENT_DIR) | |
# Run the FHE execution | |
start = time.time() | |
encrypted_output = fhe_server.run(encrypted_output, evaluation_key) | |
assert isinstance(encrypted_output, bytes) | |
fhe_execution_time = round(time.time() - start, 2) | |
# Retrieve the encrypted output path | |
encrypted_output_path = SERVER_DIR / f"{user_id}_encrypted_output" | |
# Write the file using the above path | |
with encrypted_output_path.open("wb") as f: | |
f.write(encrypted_output) | |
return JSONResponse(content=fhe_execution_time) | |
def get_output(user_id: str = Form()): | |
"""Retrieve the encrypted output from the server.""" | |
print("\nGet the output from the server ............\n") | |
# Path where the encrypted output is saved | |
encrypted_output_path = SERVER_DIR / f"{user_id}_encrypted_output" | |
# Read the file using the above path | |
with encrypted_output_path.open("rb") as f: | |
encrypted_output = f.read() | |
time.sleep(1) | |
# Send the encrypted output | |
return Response(encrypted_output) | |