Update app.py

app.py

@@ -1,8 +1,6 @@
 import gradio as gr
 from requests import head
 from transformer_vectorizer import TransformerVectorizer
-from sklearn.feature_extraction.text import TfidfVectorizer
-import numpy as np 
 from concrete.ml.deployment import FHEModelClient
 import numpy
 import os
@@ -26,18 +24,12 @@ time.sleep(5)
 # (encrypted data is too large to display in the browser)
 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"
 
 print("Loading the transformer model...")
 
 # Initialize the transformer vectorizer
 transformer_vectorizer = TransformerVectorizer()
-vectorizer = TfidfVectorizer()
 
 def clean_tmp_directory():
     # Allow 20 user keys to be stored.
@@ -57,69 +49,38 @@ def clean_tmp_directory():
         for user_id in user_ids:
             if file.name.endswith(f"{user_id}.npy"):
                 file.unlink()
-mes=[]
 
-def keygen(selected_tasks):
+
+def keygen():
     # Clean tmp directory if needed
     clean_tmp_directory()
 
     print("Initializing FHEModelClient...")
 
-
-    
-    if not selected_tasks:
-        return "choose a task first"  # 修改提示信息为英文
+    # Let's create a user_id
     user_id = numpy.random.randint(0, 2**32)
-    if "legal_rating" in selected_tasks:
-        model_names.append('legal_rating')
-                # Let's create a user_id
-
-        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()
 
-    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
-    
-        
-    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([text]).toarray()
-        if encodings.shape[1] < 1736:
-            # 在后面填充零
-            padding = np.zeros((1, 1736 - encodings.shape[1]))
-            encodings = np.hstack((encodings, padding))
-        elif encodings.shape[1] > 1736:
-            # 截取前1736列
-            encodings = encodings[:, :1736]
-    else:
-        fhe_api = FHEModelClient(FHE_MODEL_PATH, f".fhe_keys/{user_id}")
-        encodings = transformer_vectorizer.transform([text])
-    
+
+    fhe_api = FHEModelClient(FHE_MODEL_PATH, f".fhe_keys/{user_id}")
     fhe_api.load()
+    encodings = transformer_vectorizer.transform([text])
     quantized_encodings = fhe_api.model.quantize_input(encodings).astype(numpy.uint8)
     encrypted_quantized_encoding = fhe_api.quantize_encrypt_serialize(encodings)
 
@@ -137,7 +98,6 @@ def encode_quantize_encrypt(text, user_id):
     )
 
 
-
 def run_fhe(user_id):
     encoded_data_path = Path(f"tmp/tmp_encrypted_quantized_encoding_{user_id}.npy")
     if not user_id:
@@ -159,10 +119,9 @@ def run_fhe(user_id):
     query["evaluation_key"] = encoded_evaluation_key
     query["encrypted_encoding"] = encrypted_quantized_encoding
     headers = {"Content-type": "application/json"}
-
     response = requests.post(
-                "http://localhost:8000/predict_sentiment", data=json.dumps(query), headers=headers
-            )
+        "http://localhost:8000/predict_sentiment", data=json.dumps(query), headers=headers
+    )
     encrypted_prediction = base64.b64decode(response.json()["encrypted_prediction"])
 
     # Save encrypted_prediction in a file, since too large to pass through regular Gradio
@@ -183,9 +142,6 @@ 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()
 
@@ -193,12 +149,10 @@ def decrypt_prediction(user_id):
     fhe_api.generate_private_and_evaluation_keys(force=False)
 
     predictions = fhe_api.deserialize_decrypt_dequantize(encrypted_prediction)
-    print(predictions)
-    
     return {
-        "low_relative": predictions[0][0],
-        "medium_relative": predictions[0][1],
-        "high_relative": predictions[0][2],
+        "negative": predictions[0][0],
+        "neutral": predictions[0][1],
+        "positive": predictions[0][2],
     }
 
 
@@ -210,12 +164,22 @@ with demo:
 
     gr.Markdown(
         """
-
-<h2 align="center">📄Cipher Clause</h2>
-        <p align="center">
-            <img width=200 src="https://www.helloimg.com/i/2024/09/28/66f7f6701bcfb.jpeg">
-        </p>
-
+<p align="center">
+    <img width=200 src="https://user-images.githubusercontent.com/5758427/197816413-d9cddad3-ba38-4793-847d-120975e1da11.png">
+</p>
+<h2 align="center">Sentiment Analysis On Encrypted Data Using Homomorphic Encryption</h2>
+<p align="center">
+    <a href="https://github.com/zama-ai/concrete-ml"> <img style="vertical-align: middle; display:inline-block; margin-right: 3px;" width=15 src="https://user-images.githubusercontent.com/5758427/197972109-faaaff3e-10e2-4ab6-80f5-7531f7cfb08f.png">Concrete-ML</a>
+    —
+    <a href="https://docs.zama.ai/concrete-ml"> <img style="vertical-align: middle; display:inline-block; margin-right: 3px;" width=15 src="https://user-images.githubusercontent.com/5758427/197976802-fddd34c5-f59a-48d0-9bff-7ad1b00cb1fb.png">Documentation</a>
+    —
+    <a href="https://zama.ai/community"> <img style="vertical-align: middle; display:inline-block; margin-right: 3px;" width=15 src="https://user-images.githubusercontent.com/5758427/197977153-8c9c01a7-451a-4993-8e10-5a6ed5343d02.png">Community</a>
+    —
+    <a href="https://twitter.com/zama_fhe"> <img style="vertical-align: middle; display:inline-block; margin-right: 3px;" width=15 src="https://user-images.githubusercontent.com/5758427/197975044-bab9d199-e120-433b-b3be-abd73b211a54.png">@zama_fhe</a>
+</p>
+<p align="center">
+  <img src="https://user-images.githubusercontent.com/56846628/219329304-6868be9e-5ce8-4279-9123-4cb1bc0c2fb5.png" width="60%" height="60%">
+</p>
 """
     )
 
@@ -236,24 +200,10 @@ with demo:
 - The evaluation key is a public key that the server needs to process encrypted data.
 """
     )
-    gr.Markdown(
-        """
-    <hr/>
-        """
-    )
-    gr.Markdown("# Step 0: Select Task")    
-    task_checkbox = gr.CheckboxGroup(
-                choices=["legal_rating", "financial_rating"],
-                label="select_tasks"
-            )
-    gr.Markdown(
-        """
-    <hr/>
-        """
-    )
+
     gr.Markdown("# Step 1: Generate the keys")
 
-    b_gen_key_and_install = gr.Button("Generate all the keys and send public part to server")
+    b_gen_key_and_install = gr.Button("Generate the keys and send public part to server")
 
     evaluation_key = gr.Textbox(
         label="Evaluation key (truncated):",
@@ -267,46 +217,34 @@ with demo:
         interactive=False,
         visible=False
     )
-    gr.Markdown(
-        """
-<hr/>
-        """
-    )
-    gr.Markdown("# Step 2: Provide a contract or clause")
+
+    gr.Markdown("# Step 2: Provide a message")
     gr.Markdown("## Client side")
     gr.Markdown(
-        "Enter a contract or clause you want to analysis)."
-    )
-    text = gr.Textbox(label="Enter some words:", value="The Employee is entitled to two weeks of paid vacation annually, to be scheduled at the mutual convenience of the Employee and Employer.")
-    gr.Markdown(
-        """
-<hr/>
-        """
+        "Enter a sensitive text message you received and would like to do sentiment analysis on (ideas: the last text message of your boss.... or lover)."
     )
+    text = gr.Textbox(label="Enter a message:", value="I really like your work recently")
+
     gr.Markdown("# Step 3: Encode the message with the private key")
     b_encode_quantize_text = gr.Button(
-        "Encode, quantize and encrypt the text with vectorizer, and send to server"
+        "Encode, quantize and encrypt the text with transformer vectorizer, and send to server"
     )
 
     with gr.Row():
         encoding = gr.Textbox(
-            label="Representation:",
+            label="Transformer representation:",
             max_lines=4,
             interactive=False,
         )
         quantized_encoding = gr.Textbox(
-            label="Quantized  representation:", max_lines=4, interactive=False
+            label="Quantized transformer representation:", max_lines=4, interactive=False
         )
         encrypted_quantized_encoding = gr.Textbox(
-            label="Encrypted quantized representation (truncated):",
+            label="Encrypted quantized transformer representation (truncated):",
             max_lines=4,
             interactive=False,
         )
-    gr.Markdown(
-        """
-<hr/>
-        """
-    )
+
     gr.Markdown("# Step 4: Run the FHE evaluation")
     gr.Markdown("## Server side")
     gr.Markdown(
@@ -319,22 +257,18 @@ with demo:
         max_lines=4,
         interactive=False,
     )
-    gr.Markdown(
-        """
-<hr/>
-        """
-    )
-    gr.Markdown("# Step 5: Decrypt the class")
+
+    gr.Markdown("# Step 5: Decrypt the sentiment")
     gr.Markdown("## Client side")
     gr.Markdown(
         "The encrypted sentiment is sent back to client, who can finally decrypt it with its private key. Only the client is aware of the original tweet and the prediction."
     )
     b_decrypt_prediction = gr.Button("Decrypt prediction")
 
-    labels_sentiment = gr.Label(label="level:")
+    labels_sentiment = gr.Label(label="Sentiment:")
 
     # Button for key generation
-    b_gen_key_and_install.click(keygen, inputs=[task_checkbox], outputs=[evaluation_key, user_id])
+    b_gen_key_and_install.click(keygen, inputs=[], outputs=[evaluation_key, user_id])
 
     # Button to quantize and encrypt
     b_encode_quantize_text.click(