Update app.py

app.py

@@ -27,6 +27,8 @@ time.sleep(5)
 ENCRYPTED_DATA_BROWSER_LIMIT = 500
 N_USER_KEY_STORED = 20
 model_names=['financial_rating','legal_rating']
+
+
 FHE_MODEL_PATH = "deployment/financial_rating"
 FHE_LEGAL_PATH = "deployment/legal_rating"
 #FHE_LEGAL_PATH="deployment/legal_rating"
@@ -55,7 +57,7 @@ def clean_tmp_directory():
         for user_id in user_ids:
             if file.name.endswith(f"{user_id}.npy"):
                 file.unlink()
-model_nams=[]
+model_names=[]
 
 def keygen(selected_tasks):
     # Clean tmp directory if needed
@@ -63,36 +65,47 @@ def keygen(selected_tasks):
 
     print("Initializing FHEModelClient...")
 
+
     
     if not selected_tasks:
-        return "choose task first"  # 修改提示信息为英文
+        return "choose a task first"  # 修改提示信息为英文
 
     if "legal_rating" in selected_tasks:
         model_names.append('legal_rating')
-    if "financial_rating" in selected_tasks:
-        model_names.append('financial_rating')
+                # Let's create a user_id
 
+        fhe_api= FHEModelClient(FHE_LEGAL_PATH, f".fhe_keys/{user_id}")
 
 
-    # Let's create a user_id
+    if "financial_rating" in selected_tasks:
+        model_names.append('financial_rating')
+
+        fhe_api = FHEModelClient(FHE_MODEL_PATH, f".fhe_keys/{user_id}")
+        
+        # Let's create a user_id
     user_id = numpy.random.randint(0, 2**32)
-    fhe_api = FHEModelClient(FHE_MODEL_PATH, f".fhe_keys/{user_id}")
+        
     fhe_api.load()
-
-
-    # Generate a fresh key
+    
+        
+        # Generate a fresh key
     fhe_api.generate_private_and_evaluation_keys(force=True)
     evaluation_key = fhe_api.get_serialized_evaluation_keys()
-
-    # Save evaluation_key in a file, since too large to pass through regular Gradio
-    # buttons, https://github.com/gradio-app/gradio/issues/1877
+    
+        # Save evaluation_key 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_evaluation_key_{user_id}.npy", evaluation_key)
 
-    return [list(evaluation_key)[:ENCRYPTED_DATA_BROWSER_LIMIT], user_id]
+
+
+
+    
 def encode_quantize_encrypt(text, user_id):
     if not user_id:
         raise gr.Error("You need to generate FHE keys first.")
-
+    if "legal_rating" in model_names:
+        fhe_api = FHEModelClient(FHE_LEGAL_PATH, f".fhe_keys/{user_id}")
+        encodings =vectorizer.fit_transform(data['clause'])
     fhe_api = FHEModelClient(FHE_MODEL_PATH, f".fhe_keys/{user_id}")
     fhe_api.load()
     encodings = transformer_vectorizer.transform([text])
@@ -158,6 +171,9 @@ def decrypt_prediction(user_id):
     # Read encrypted_prediction from the file
     encrypted_prediction = numpy.load(encoded_data_path).tobytes()
 
+    if "legal_rating" in model_names:
+        fhe_api = FHEModelClient(FHE_LEGAL_PATH, f".fhe_keys/{user_id}")
+
     fhe_api = FHEModelClient(FHE_MODEL_PATH, f".fhe_keys/{user_id}")
     fhe_api.load()
 
@@ -166,6 +182,7 @@ def decrypt_prediction(user_id):
 
     predictions = fhe_api.deserialize_decrypt_dequantize(encrypted_prediction)
     print(predictions)
+    
     return {
         "low_relative": predictions[0][0],
         "medium_relative": predictions[0][1],
@@ -204,7 +221,7 @@ with demo:
 - The evaluation key is a public key that the server needs to process encrypted data.
 """
     )
-    gr.Markdown("# Step 0: Select Tasks")    
+    gr.Markdown("# Step 0: Select Task")    
     task_checkbox = gr.CheckboxGroup(
                 choices=["legal_rating", "financial_rating"],
                 label="select_tasks"