Upload app.py

app.py

@@ -18,7 +18,7 @@ import time
 import pandas as pd
 import pickle
 import numpy as np
-
+import pdb
 # This repository's directory
 REPO_DIR = Path(__file__).parent
 subprocess.Popen(["uvicorn", "server:app"], cwd=REPO_DIR)
@@ -82,38 +82,50 @@ def keygen():
     return [list(evaluation_key)[:ENCRYPTED_DATA_BROWSER_LIMIT], user_id]
 
 
-def encode_quantize_encrypt(test_file, user_id):
-
-    fhe_api = FHEModelClient(f"fhe_model", f".fhe_keys/{user_id}")
+def encode_quantize_encrypt(test_file, eval_key):
+    ugly = ['Machine', 'SizeOfOptionalHeader', 'Characteristics',
+       'MajorLinkerVersion', 'MinorLinkerVersion', 'SizeOfCode',
+       'SizeOfInitializedData', 'SizeOfUninitializedData',
+       'AddressOfEntryPoint', 'BaseOfCode', 'BaseOfData', 'ImageBase',
+       'SectionAlignment', 'FileAlignment', 'MajorOperatingSystemVersion',
+       'MinorOperatingSystemVersion', 'MajorImageVersion', 'MinorImageVersion',
+       'MajorSubsystemVersion', 'MinorSubsystemVersion', 'SizeOfImage',
+       'SizeOfHeaders', 'CheckSum', 'Subsystem', 'DllCharacteristics',
+       'SizeOfStackReserve', 'SizeOfStackCommit', 'SizeOfHeapReserve',
+       'SizeOfHeapCommit', 'LoaderFlags', 'NumberOfRvaAndSizes', 'SectionsNb',
+       'SectionsMeanEntropy', 'SectionsMinEntropy', 'SectionsMaxEntropy',
+       'SectionsMeanRawsize', 'SectionsMinRawsize',
+       'SectionsMeanVirtualsize', 'SectionsMinVirtualsize',
+       'SectionMaxVirtualsize', 'ImportsNbDLL', 'ImportsNb',
+       'ImportsNbOrdinal', 'ExportNb', 'ResourcesNb', 'ResourcesMeanEntropy',
+       'ResourcesMinEntropy', 'ResourcesMaxEntropy', 'ResourcesMeanSize',
+       'ResourcesMinSize', 'ResourcesMaxSize', 'LoadConfigurationSize',
+       'VersionInformationSize']
+
+    fhe_api = FHEModelClient(f"fhe_model", f".fhe_keys/{eval_key}")
     fhe_api.load()
     from PE_main import extract_infos
-
-    features = pickle.loads(open(os.path.join("features.pkl"), "rb").read())
+    # expect [1, 53] but we get (53)
+    # pdb.set_trace()
+    # features = pickle.loads(open(os.path.join("features.pkl"), "rb").read())
     encodings = extract_infos(test_file)
-    encodings = list(map(lambda x: encodings[x], features))
+
+    encodings = list(map(lambda x: encodings[x], ugly))
+
+    encodings = np.array(encodings).reshape(1, -1)
 
     quantized_encodings = fhe_api.model.quantize_input(encodings).astype(numpy.uint8)
     encrypted_quantized_encoding = fhe_api.quantize_encrypt_serialize(encodings)
 
-    # Save encrypted_quantized_encoding in a file, since too large to pass through regular Gradio
-    # buttons, https://github.com/gradio-app/gradio/issues/1877
     numpy.save(
-        f"tmp/tmp_encrypted_quantized_encoding_{user_id}.npy",
+        f"tmp/tmp_encrypted_quantized_encoding_{eval_key[1]}.npy",
         encrypted_quantized_encoding,
     )
 
     # Compute size
-    encrypted_quantized_encoding_shorten = list(encrypted_quantized_encoding)[
-        :ENCRYPTED_DATA_BROWSER_LIMIT
-    ]
-    encrypted_quantized_encoding_shorten_hex = "".join(
-        f"{i:02x}" for i in encrypted_quantized_encoding_shorten
-    )
-    return (
-        encodings[0],
-        quantized_encodings[0],
-        encrypted_quantized_encoding_shorten_hex,
-    )
+    encrypted_quantized_encoding_shorten = list(encrypted_quantized_encoding)[:ENCRYPTED_DATA_BROWSER_LIMIT]
+    encrypted_quantized_encoding_shorten_hex = "".join(f"{i:02x}" for i in encrypted_quantized_encoding_shorten)
+    return (encodings[0],quantized_encodings[0],encrypted_quantized_encoding_shorten_hex)
 
 
 def run_fhe(user_id):
@@ -124,9 +136,8 @@ def run_fhe(user_id):
     evaluation_key = numpy.load(f"tmp/tmp_evaluation_key_{user_id}.npy")
 
     # Use base64 to encode the encodings and evaluation key
-    encrypted_quantized_encoding = base64.b64encode(
-        encrypted_quantized_encoding
-    ).decode()
+    encrypted_quantized_encoding = base64.b64encode(encrypted_quantized_encoding).decode()
+
     encoded_evaluation_key = base64.b64encode(evaluation_key).decode()
 
     query = {}
@@ -134,23 +145,19 @@ def run_fhe(user_id):
     query["encrypted_encoding"] = encrypted_quantized_encoding
     headers = {"Content-type": "application/json"}
 
-
     response = requests.post(
         "http://localhost:8000/predict",
         data=json.dumps(query),
         headers=headers,
     )
 
-
     encrypted_prediction = base64.b64decode(response.json()["encrypted_prediction"])
 
     numpy.save(f"tmp/tmp_encrypted_prediction_{user_id}.npy", encrypted_prediction)
-    encrypted_prediction_shorten = list(encrypted_prediction)[
-        :ENCRYPTED_DATA_BROWSER_LIMIT
-    ]
-    encrypted_prediction_shorten_hex = "".join(
-        f"{i:02x}" for i in encrypted_prediction_shorten
-    )
+
+    encrypted_prediction_shorten = list(encrypted_prediction)[:ENCRYPTED_DATA_BROWSER_LIMIT]
+    encrypted_prediction_shorten_hex = "".join(f"{i:02x}" for i in encrypted_prediction_shorten)
+    return encrypted_prediction_shorten_hex
 
 
 def decrypt_prediction(user_id):
@@ -167,20 +174,23 @@ def decrypt_prediction(user_id):
     fhe_api.generate_private_and_evaluation_keys(force=False)
 
     predictions = fhe_api.deserialize_decrypt_dequantize(encrypted_prediction)
+    return predictions
 
 
 
 def process_pipeline(test_file):
+
     eval_key = keygen()
     encodings = encode_quantize_encrypt(test_file, eval_key)
-    encrypted_quantized_encoding = run_fhe(encodings)
-    encrypted_prediction = decrypt_prediction(encrypted_quantized_encoding)
+    encrypted_quantized_encoding = run_fhe(eval_key[1])
+    encrypted_prediction = decrypt_prediction(eval_key[1])
 
     return eval_key, encodings, encrypted_quantized_encoding, encrypted_prediction
 
 if __name__ == "__main__":
-
+        
     app = gr.Interface(
+
         fn=process_pipeline,
         inputs=[
             gr.File(label="Test File"),