Update app.py

app.py

@@ -521,7 +521,6 @@ with demo:
     """
     <p align="center" style="font-size: 16px;">
         Anonymization makes it impossible to identify a person from a data set and thus allows their privacy to be respected </p> Concealing some of PII information is a helpful technique for de-identification. However, this can't be considered as anonymization. 
-
     <p align="center" style="font-size: 16px;">
         De-identification uses Fully Homomorphic Encryption (FHE) to conceal personally 
         identifiable information (PII) within encrypted documents, enabling computations to be 
@@ -561,12 +560,9 @@ with demo:
     with gr.Row():
         audio_recorder = gr.Audio(source="microphone", type="file", label="Record Audio")
         record_output = gr.Audio(label="Recorded Audio", type="filepath")
-
     # When the user records an audio, save it
     audio_recorder.change(fn=save_recorded_audio, inputs=[audio_recorder], outputs=[record_output])
-
     gen_key_btn = gr.Button("Generate the secret and evaluation keys")
-
     gen_key_btn.click(
         key_gen_fn,
         inputs=[],
@@ -617,32 +613,32 @@ with demo:
         outputs=[gen_key_btn],
     )
 
-    ########################### Main document Part ##########################
+    ########################## Main document Part ##########################
 
-   # gr.Markdown("<hr />")
-   # gr.Markdown("## Step 2.1: Select the document you want to encrypt\n\n"
-    #    """To make it simple, we pre-compiled the following document, but you are free to choose 
-    #    on which part you want to run this example.
-   #     """
-#    )
+    gr.Markdown("<hr />")
+    gr.Markdown("## Step 2.1: Select the document you want to encrypt\n\n"
+        """To make it simple, we pre-compiled the following document, but you are free to choose 
+        on which part you want to run this example.
+        """
+    )
 
-   # with gr.Row():
-   #     with gr.Column(scale=5):
-   #         original_sentences_box = gr.CheckboxGroup(
-   #             ORIGINAL_DOCUMENT,
-   #             value=ORIGINAL_DOCUMENT,
-    #            label="Contract:",
-   #             show_label=True,
-    #        )
+    with gr.Row():
+        with gr.Column(scale=5):
+            original_sentences_box = gr.CheckboxGroup(
+                ORIGINAL_DOCUMENT,
+                value=ORIGINAL_DOCUMENT,
+                label="Contract:",
+                show_label=True,
+            )
 
-   #     with gr.Column(scale=1, min_width=6):
-    #        gr.HTML("<div style='height: 77px;'></div>")
-    #        encrypt_doc_btn = gr.Button("Encrypt the document")
+        with gr.Column(scale=1, min_width=6):
+            gr.HTML("<div style='height: 77px;'></div>")
+            encrypt_doc_btn = gr.Button("Encrypt the document")
 
-    #    with gr.Column(scale=5):
-    #        encrypted_doc_box = gr.Textbox(
-    #            label="Encrypted document:", show_label=True, interactive=False, lines=10
-    #        )
+        with gr.Column(scale=5):
+            encrypted_doc_box = gr.Textbox(
+                label="Encrypted document:", show_label=True, interactive=False, lines=10
+            )
 
 
     ########################## User Query Part ##########################
@@ -659,13 +655,13 @@ with demo:
 
             with gr.Column(scale=5):
                 default_query_box = gr.Dropdown(
-                    list(DEFAULT_QUERIES.values()), label="PROMPT"
+                    list(DEFAULT_QUERIES.values()), label="PROMPT EXAMPLES:"
                 )
 
             gr.Markdown("Or")
 
             query_box = gr.Textbox(
-                value=transcription_output, label="CUSTOMIZED PROMPT:", interactive=True
+                value=" Hello. My name is Inuitvementoya. You kill my father. Prepare to die.", label="CUSTOMIZED PROMPT:", interactive=True
             )
 
             default_query_box.change(
@@ -686,24 +682,40 @@ with demo:
             )
 
     ########################## FHE processing Part ##########################
+
     gr.Markdown("<hr />")
-    gr.Markdown("## Step 3: De-identify the prompt using FHE")
+    gr.Markdown("## Step 3: De-identify the document and the prompt using FHE")
     gr.Markdown(
-        """Once the client encrypts the prompt locally, it will be sent to a remote 
+        """Once the client encrypts the document and the prompt locally, it will be sent to a remote 
         server to perform the de-identification on encrypted data. When the computation is done, the 
         server will return the result to the client for decryption."""
+        
     )
 
     run_fhe_btn = gr.Button("De-identify using FHE")
 
     with gr.Row():
         with gr.Column(scale=5):
+
+            anonymized_doc_output = gr.Textbox(
+                label="Decrypted and de-idenntified document", lines=10, interactive=True
+            )
+
+        with gr.Column(scale=5):
+
             anonymized_query_output = gr.Textbox(
                 label="Decrypted and de-identified prompt", lines=10, interactive=True
             )
-    
+
+
     identified_words_output_df = gr.Dataframe(label="Identified words:", visible=False)
 
+    encrypt_doc_btn.click(
+        fn=encrypt_doc_fn,
+        inputs=[original_sentences_box],
+        outputs=[encrypted_doc_box, anonymized_doc_output],
+    )
+
     encrypt_query_btn.click(
         fn=encrypt_query_fn,
         inputs=[query_box],
@@ -717,9 +729,10 @@ with demo:
 
     run_fhe_btn.click(
         anonymization_with_fn,
-        inputs=[query_box],
-        outputs=[anonymized_query_output, identified_words_output_df],
+        inputs=[original_sentences_box, query_box],
+        outputs=[anonymized_doc_output, anonymized_query_output, identified_words_output_df],
     )
+    
 
     ########################## Presidio ##########################
     gr.Markdown("<hr />")
@@ -744,4 +757,4 @@ with demo:
 
 
 # Launch the app
-demo.launch(share=False)
+demo.launch(share=False)