File size: 2,201 Bytes
be526a2 df9cd1e be526a2 8675aad be526a2 25daa45 be526a2 25daa45 be526a2 145adb7 be526a2 145adb7 1977567 145adb7 90e521b 145adb7 e06e968 90e521b 145adb7 b16eb1c 145adb7 1bb01e9 46dab14 145adb7 46dab14 78898ff a59f4d1 145adb7 46dab14 be526a2 46dab14 be526a2 46dab14 be526a2 46dab14 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 |
from typing import Dict, List, Any
import numpy as np
from concrete.ml.deployment import FHEModelServer
def from_json(python_object):
if "__class__" in python_object:
return bytes(python_object["__value__"])
def to_json(python_object):
if isinstance(python_object, bytes):
return {"__class__": "bytes", "__value__": list(python_object)}
raise TypeError(repr(python_object) + " is not JSON serializable")
class EndpointHandler:
def __init__(self, path=""):
# For server
self.fhemodel_server = FHEModelServer(path + "/compiled_model")
# Simulate a database of keys
self.key_database = {}
def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
"""
data args:
inputs (:obj: `str`)
date (:obj: `str`)
Return:
A :obj:`list` | `dict`: will be serialized and returned
"""
# Get method
method = data.pop("method", data)
if method == "save_key":
# Get keys
evaluation_keys = from_json(data.pop("evaluation_keys", data))
uid = np.random.randint(2**32)
while uid in self.key_database.keys():
uid = np.random.randint(2**32)
self.key_database[uid] = evaluation_keys
return {"uid": uid}
elif method == "append_key":
# Get key piece
evaluation_keys = from_json(data.pop("evaluation_keys", data))
uid = data.pop("uid", data)
self.key_database[uid] += evaluation_keys
return
elif method == "inference":
uid = data.pop("uid", data)
assert uid in self.key_database.keys(), f"{uid} not in DB, {self.key_database.keys()=}"
# Get inputs
encrypted_inputs = from_json(data.pop("encrypted_inputs", data))
# Find key in the database
evaluation_keys = self.key_database[uid]
# Run CML prediction
encrypted_prediction = self.fhemodel_server.run(encrypted_inputs, evaluation_keys)
return to_json(encrypted_prediction)
else:
return
|