chore: add key gen

app.py

@@ -9,6 +9,7 @@ import pandas as pd
 from fhe_anonymizer import FHEAnonymizer
 from openai import OpenAI
 from utils_demo import *
+from concrete.ml.deployment import FHEModelClient
 
 ORIGINAL_DOCUMENT = read_txt(ORIGINAL_FILE_PATH).split("\n\n")
 ANONYMIZED_DOCUMENT = read_txt(ANONYMIZED_FILE_PATH)
@@ -40,9 +41,27 @@ def key_gen_fn() -> Dict:
     """
     print("Key Gen..")
 
-    anonymizer.generate_key()
+    # Generate a random user ID
+    user_id = np.random.randint(0, 2**32)
+    print(f"Your user ID is: {user_id}....")
 
-    evaluation_key_path = KEYS_DIR / "evaluation_key"
+    client = FHEModelClient(path_dir=DEPLOYMENT_DIR, key_dir=KEYS_DIR / f"{user_id}")
+    client.load()
+
+    # Creates the private and evaluation keys on the client side
+    client.generate_private_and_evaluation_keys()
+
+    # Get the serialized evaluation keys
+    serialized_evaluation_keys = client.get_serialized_evaluation_keys()
+    assert isinstance(serialized_evaluation_keys, bytes)
+
+    # Save the evaluation key
+    evaluation_key_path = KEYS_DIR / f"{user_id}/evaluation_key"
+    
+    with evaluation_key_path.open("wb") as f:
+        f.write(serialized_evaluation_keys)
+
+    # anonymizer.generate_key()
 
     if not evaluation_key_path.is_file():
         error_message = (
@@ -244,8 +263,6 @@ with demo:
 
     gr.Markdown("## Step 2: Private document")
 
-
-
     with gr.Row():
         with gr.Column():
             gr.Markdown("**Original document:**")

utils_demo.py

@@ -25,7 +25,7 @@ CURRENT_DIR = Path(__file__).parent
 DATA_PATH = CURRENT_DIR / "files"
 LOGREG_MODEL_PATH = CURRENT_DIR / "models" / "cml_logreg.model"
 DEPLOYMENT_DIR = CURRENT_DIR / "deployment"
-KEYS_DIR = DEPLOYMENT_DIR / ".fhe_keys"
+KEYS_DIR = DEPLOYMENT_DIR / "fhe_keys"
 
 ORIGINAL_FILE_PATH = DATA_PATH / "original_document.txt"
 ANONYMIZED_FILE_PATH = DATA_PATH / "anonymized_document.txt"