Update app.py

app.py

@@ -105,7 +105,29 @@ def create_mask_dict(entities):
                     entity_counters[entity['entity_group']] += 1
                 mask_dict[entity['word']] = f"{entity['entity_group']}_{entity_counters[entity['entity_group']]}"
     return mask_dict
-
+def export_masked_text(masked_text, file_type):
+    if file_type == "txt":
+        return masked_text.encode("utf-8")
+    elif file_type == "pdf":
+        pdf_buffer = io.BytesIO()
+        from fpdf import FPDF
+        pdf = FPDF()
+        pdf.add_page()
+        pdf.set_font("Arial", size=12)
+        pdf.multi_cell(0, 10, masked_text)
+        pdf.output(pdf_buffer)
+        pdf_buffer.seek(0)
+        return pdf_buffer.getvalue()
+    elif file_type == "docx":
+        doc = docx.Document()
+        doc.add_paragraph(masked_text)
+        buffer = io.BytesIO()
+        doc.save(buffer)
+        buffer.seek(0)
+        return buffer.getvalue()
+    else:
+        st.error("Unsupported file type for export")
+        return None
 Run_Button = st.button("Run")
 
 if Run_Button and input_text:
@@ -147,6 +169,8 @@ if Run_Button and input_text:
     
     st.subheader("Recognized Entities")
     st.dataframe(df_final)
+
+    
     
     # Spacy display logic with entity numbering
     spacy_display = {"ents": [], "text": input_text, "title": None}
@@ -160,5 +184,16 @@ if Run_Button and input_text:
     html = spacy.displacy.render(spacy_display, style="ent", minify=True, manual=True)
     st.write(html, unsafe_allow_html=True)
 
+    export_file_type = uploaded_file.type.split("/")[-1] if uploaded_file is not None else "txt"
+    if st.button("Download Masked Text"):
+        masked_file_content = export_masked_text(masked_text, export_file_type)
+        if masked_file_content:
+            st.download_button(
+                label="Download",
+                data=masked_file_content,
+                file_name=f"masked_output.{export_file_type}",
+                mime=f"application/{export_file_type}" if export_file_type != "txt" else "text/plain"
+            )
+
     st.subheader("Masking Dictionary")
     st.json(mask_dict)