File size: 1,943 Bytes
be526a2
 
df9cd1e
be526a2
 
8675aad
 
 
 
 
 
 
 
 
 
 
be526a2
25daa45
be526a2
 
25daa45
be526a2
145adb7
 
 
be526a2
 
 
 
 
 
 
 
 
145adb7
 
 
1977567
145adb7
 
 
 
90e521b
145adb7
e06e968
90e521b
145adb7
 
 
b16eb1c
145adb7
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
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 == "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