chore: clean a bit

play_with_endpoint.py

@@ -50,7 +50,7 @@ _, X_test, _, Y_test = train_test_split(
     random_state=42,
 )
 
-NB_SAMPLES = 2
+NB_SAMPLES = 10
 X_test = X_test[:NB_SAMPLES]
 Y_test = Y_test[:NB_SAMPLES]
 
@@ -67,15 +67,15 @@ nb_samples = len(X_test)
 verbose = False
 time_start = time.time()
 duration = 0
-is_first = True
 
 for i in range(nb_samples):
 
     # Quantize the input and encrypt it
     encrypted_inputs = fhemodel_client.quantize_encrypt_serialize(X_test[i].reshape(1, -1))
 
-    print(f"Size of encrypted input {sys.getsizeof(encrypted_inputs)}")
-    print(f"Size of keys {sys.getsizeof(evaluation_keys)}")
+    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 = {
@@ -84,32 +84,22 @@ for i in range(nb_samples):
         "evaluation_keys": to_json(evaluation_keys),
     }
 
-    print(f"{payload=}")
-
     # Run the inference on HF servers
     duration -= time.time()
-    print(f"Starting at {time.time()}")
     encrypted_prediction = query(payload)
-    print(f"Ending at {time.time()}")
     duration += time.time()
 
-    print(f"{encrypted_prediction=}")
-
-    encrypted_prediction = encrypted_prediction
-
-    if is_first:
-        is_first = False
-        print(f"Size of the payload: {sys.getsizeof(payload)} bytes")
+    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]}")
 
     # Measure accuracy
-    nb_good += y_test[i] == prediction
+    nb_good += Y_test[i] == prediction
 
 print(f"Accuracy on {nb_samples} samples is {nb_good * 1. / nb_samples}")
 print(f"Total time: {time.time() - time_start} seconds")