|
"""Modified classes for use for Client-Server interface with multi-inputs circuits.""" |
|
|
|
import numpy |
|
import copy |
|
|
|
from concrete.fhe import Value, EvaluationKeys |
|
|
|
from concrete.ml.deployment.fhe_client_server import FHEModelClient, FHEModelDev, FHEModelServer |
|
from concrete.ml.sklearn import XGBClassifier as ConcreteXGBClassifier |
|
|
|
|
|
class MultiInputsFHEModelDev(FHEModelDev): |
|
|
|
def __init__(self, *arg, **kwargs): |
|
|
|
super().__init__(*arg, **kwargs) |
|
|
|
model = copy.copy(self.model) |
|
model.__class__ = ConcreteXGBClassifier |
|
self.model = model |
|
|
|
|
|
class MultiInputsFHEModelClient(FHEModelClient): |
|
|
|
def __init__(self, *args, nb_inputs=1, **kwargs): |
|
self.nb_inputs = nb_inputs |
|
|
|
super().__init__(*args, **kwargs) |
|
|
|
def quantize_encrypt_serialize_multi_inputs( |
|
self, |
|
x: numpy.ndarray, |
|
input_index, |
|
processed_input_shape, |
|
input_slice |
|
) -> bytes: |
|
|
|
x_padded = numpy.zeros(processed_input_shape) |
|
|
|
x_padded[:, input_slice] = x |
|
|
|
q_x_padded = self.model.quantize_input(x_padded) |
|
|
|
q_x = q_x_padded[:, input_slice] |
|
|
|
q_x_inputs = [None for _ in range(self.nb_inputs)] |
|
q_x_inputs[input_index] = q_x |
|
|
|
|
|
q_x_enc = self.client.encrypt(*q_x_inputs) |
|
|
|
|
|
q_x_enc_ser = q_x_enc[input_index].serialize() |
|
return q_x_enc_ser |
|
|
|
|
|
class MultiInputsFHEModelServer(FHEModelServer): |
|
|
|
def run( |
|
self, |
|
*serialized_encrypted_quantized_data: bytes, |
|
serialized_evaluation_keys: bytes, |
|
) -> bytes: |
|
"""Run the model on the server over encrypted data. |
|
|
|
Args: |
|
serialized_encrypted_quantized_data (bytes): the encrypted, quantized |
|
and serialized data |
|
serialized_evaluation_keys (bytes): the serialized evaluation keys |
|
|
|
Returns: |
|
bytes: the result of the model |
|
""" |
|
assert self.server is not None, "Model has not been loaded." |
|
|
|
deserialized_encrypted_quantized_data = tuple(Value.deserialize(data) for data in serialized_encrypted_quantized_data) |
|
|
|
deserialized_evaluation_keys = EvaluationKeys.deserialize(serialized_evaluation_keys) |
|
|
|
result = self.server.run( |
|
*deserialized_encrypted_quantized_data, evaluation_keys=deserialized_evaluation_keys |
|
) |
|
serialized_result = result.serialize() |
|
return serialized_result |