chore: restore json

handler.py

@@ -3,6 +3,17 @@ 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=""):
 
@@ -19,12 +30,12 @@ class EndpointHandler:
         """
 
         # Get inputs
-        encrypted_inputs = data.pop("encrypted_inputs", data)
+        encrypted_inputs = from_json(data.pop("encrypted_inputs", data))
 
         # Get keys
-        evaluation_keys = data.pop("evaluation_keys", data)
+        evaluation_keys = from_json(data.pop("evaluation_keys", data))
 
         # Run CML prediction
         encrypted_prediction = self.fhemodel_server.run(encrypted_inputs, evaluation_keys)
 
-        return encrypted_prediction
+        return to_json(encrypted_prediction)

play_with_endpoint.py

@@ -12,15 +12,26 @@ from concrete.ml.deployment import FHEModelClient
 import requests
 
 
+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")
+
+
+def from_json(python_object):
+    if "__class__" in python_object:
+        return bytes(python_object["__value__"])
+
+
 API_URL = "https://puqif7goarh132kl.us-east-1.aws.endpoints.huggingface.cloud"
 headers = {
     "Authorization": "Bearer " + os.environ.get("HF_TOKEN"),
-    "Content-Type": "application/octet-stream",
+    "Content-Type": "application/json",
 }
 
 
 def query(payload):
-    response = requests.post(API_URL, headers=headers, data=payload)
+    response = requests.post(API_URL, headers=headers, json=payload)
     return response.json()
 
 
@@ -47,26 +58,34 @@ for i in range(nb_samples):
     # Quantize the input and encrypt it
     encrypted_inputs = fhemodel_client.quantize_encrypt_serialize([X_test[i]])
 
+    if verbose:
+        print(f"Size of encrypted input: {sys.getsizeof(encrypted_inputs) / 1024 / 1024} megabytes")
+        print(f"Size of keys: {sys.getsizeof(evaluation_keys) / 1024 / 1024} megabytes")
+
     # Prepare the payload, including the evaluation keys which are needed server side
     payload = {
         "inputs": "fake",
-        "encrypted_inputs": encrypted_inputs,
-        "evaluation_keys": evaluation_keys,
+        "encrypted_inputs": to_json(encrypted_inputs),
+        "evaluation_keys": to_json(evaluation_keys),
     }
 
     # Run the inference on HF servers
     duration -= time.time()
+    duration_inference = -time.time()
     encrypted_prediction = query(payload)
     duration += time.time()
+    duration_inference += time.time()
 
-    encrypted_prediction = encrypted_prediction
+    encrypted_prediction = from_json(encrypted_prediction)
 
     # Decrypt the result and dequantize
     prediction_proba = fhemodel_client.deserialize_decrypt_dequantize(encrypted_prediction)[0]
     prediction = np.argmax(prediction_proba)
 
     if verbose or True:
-        print(f"for {i}-th input, {prediction=} with expected {y_test[i]}")
+        print(
+            f"for {i}-th input, {prediction=} with expected {y_test[i]} in {duration_inference} seconds"
+        )
 
     # Measure accuracy
     nb_good += y_test[i] == prediction