Upload app.py

app.py

@@ -25,8 +25,9 @@ PII_ENTITIES = [
 ]
 
 REDACTION_METHODS = {
+    "Remove": OperatorConfig("redact"),
     "Redact": OperatorConfig("replace", {"new_value": "[REDACTED]"}),
-    "Remove": OperatorConfig("remove"),
+    "Replace": OperatorConfig("replace", {"new_value": ""}),
     "Mask": OperatorConfig("mask", {"chars_to_mask": 4, "masking_char": "*", "from_end": False}),
     "Hash": "hash",
     "Encrypt": "encrypt",
@@ -39,7 +40,8 @@ def encrypt_pii(text):
     return fernet.encrypt(text.encode("utf-8")).decode("utf-8")
 
 def redact_text(text, selected_entities, redaction_method):
-    selected_entities = selected_entities or None
+    """Identifies and redacts selected PII types based on the chosen method."""
+    selected_entities = selected_entities or None  # If empty, redact all entities
     results = analyzer.analyze(text=text, entities=selected_entities, language="en")
     
     if redaction_method == "Hash":
@@ -48,6 +50,9 @@ def redact_text(text, selected_entities, redaction_method):
     elif redaction_method == "Encrypt":
         for result in results:
             text = text.replace(result.text, encrypt_pii(result.text))
+    elif redaction_method == "Replace":
+        operators = {entity.entity_type: OperatorConfig("replace") for entity in results}
+        text = anonymizer.anonymize(text=text, analyzer_results=results, operators=operators).text
     else:
         operators = {entity.entity_type: REDACTION_METHODS[redaction_method] for entity in results}
         text = anonymizer.anonymize(text=text, analyzer_results=results, operators=operators).text
@@ -84,28 +89,36 @@ def read_document(file):
         return process_txt(file)
 
 def save_redacted_file(original_file, redacted_text):
+    """Saves redacted text back in the original format and returns a valid file path."""
+    
     ext = original_file.name.split(".")[-1].lower()
-    temp_dir = tempfile.gettempdir()
+    
+    # Create a temporary file with a proper path
+    temp_dir = tempfile.gettempdir()  # OS-specific temp folder
     safe_filename = f"redacted_{os.path.basename(original_file.name)}"
     redacted_file_path = os.path.join(temp_dir, safe_filename)
-    
-    if ext == "docx":
-        doc = Document(original_file.name)
-        for para in doc.paragraphs:
-            para.text = redact_text(para.text, PII_ENTITIES, "Redact")
-        doc.save(redacted_file_path)
+
+    if ext == "pdf":
+        # Assuming a simple text-based redacted PDF (can be further improved for real PDF processing)
+        with open(redacted_file_path, "w", encoding="utf-8") as f:
+            f.write(redacted_text)  # Saving as text-based PDF
+    elif ext == "docx":
+        doc = Document()
+        for line in redacted_text.split("\n"):
+            doc.add_paragraph(line)
+        doc.save(redacted_file_path)  # Save DOCX in temp folder
     elif ext == "pptx":
-        ppt = Presentation(original_file.name)
-        for slide in ppt.slides:
-            for shape in slide.shapes:
-                if hasattr(shape, "text"):
-                    shape.text = redact_text(shape.text, PII_ENTITIES, "Redact")
+        ppt = Presentation()
+        slide = ppt.slides.add_slide(ppt.slide_layouts[5])  # Blank slide
+        textbox = slide.shapes.add_textbox(0, 0, ppt.slide_width, ppt.slide_height)
+        textbox.text = redacted_text
         ppt.save(redacted_file_path)
-    else:
+    else:  # TXT
         with open(redacted_file_path, "w", encoding="utf-8") as f:
             f.write(redacted_text)
-    
-    return redacted_file_path
+
+    return redacted_file_path  # Return the safe file path
+
 
 def process_file(file, selected_entities, redaction_method):
     text = read_document(file)
@@ -122,7 +135,7 @@ def deselect_all_entities():
 custom_css = """
 <style>
     #redact_button {
-        background-color: #E691FF !important;
+        /*background-color: #E691FF !important;*/
         color: #4B23C0;
     }
 </style>
@@ -130,30 +143,63 @@ custom_css = """
 
 # Gradio UI
 with gr.Blocks() as app:
+    
     gr.Markdown(
         """
-        <div style="background-color: #4B23C0; color: white; padding: 20px; text-align: left; font-size: 24px; font-weight: bold; margin: 0; border-radius: 4px;">
+        <div style="
+            background-color: #4B23C0; 
+            color: white; 
+            padding: 20px; 
+            text-align: left; 
+            font-size: 24px; 
+            font-weight: bold; 
+            margin: 0;
+            border-radius: 4px; /* Rounded edges */
+        ">
             🔒 PII Remover &nbsp;-&nbsp; Secure Document Redaction Tool
         </div>
         """,
         sanitize_html=False
     )
+    
+    gr.Markdown(
+        "<div style='text-align: center; font-size: 24px; font-weight: bold; color: red;'>"
+        "⚠️ THIS IS A DEMONSTRATION. DO NOT UPLOAD SENSITIVE DOCUMENTS. ⚠️"
+        "</div>",
+        sanitize_html=False
+    )
+
+
     gr.Markdown("Upload a **TXT, DOCX, PPTX, or PDF** file to remove **Personal Identifiable Information (PII)** while keeping formatting.")
+
+    # Load CSS
     gr.HTML(custom_css)
     
     with gr.Row():
         file_input = gr.File(label="Upload Document (PDF, DOCX, PPTX, TXT)")
-    entity_selector = gr.CheckboxGroup(PII_ENTITIES, label="Select PII Entities to Redact")
+    
+    entity_selector = gr.CheckboxGroup(PII_ENTITIES, label="Select PII Entities to Redact (Leave blank to redact all)")
+    
     with gr.Row():
         select_all_button = gr.Button("Select All")
         deselect_all_button = gr.Button("Deselect All")
-    redaction_method = gr.Radio(["Redact", "Remove", "Mask", "Hash", "Encrypt"], label="Redaction Method", value="Redact")
+    
+    redaction_method = gr.Radio(
+        ["Remove", "Redact", "Replace", "Mask", "Hash", "Encrypt"], 
+        label="Redaction Method", 
+        value="Redact"
+    )
+    
     process_button = gr.Button("Redact Document", elem_id="redact_button")
+    
     output_text = gr.Textbox(label="Redacted Text", lines=10)
     download_button = gr.File(label="Download Redacted File")
-    
+
+    # Button Actions
     select_all_button.click(fn=select_all_entities, outputs=entity_selector)
     deselect_all_button.click(fn=deselect_all_entities, outputs=entity_selector)
-    process_button.click(fn=process_file, inputs=[file_input, entity_selector, redaction_method], outputs=[output_text, download_button])
     
+    process_button.click(fn=process_file, inputs=[file_input, entity_selector, redaction_method], outputs=[output_text, download_button])
+
 app.launch()
+