Update app.py

app.py

@@ -105,12 +105,18 @@ 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 create_masked_text(input_text, entities, mask_dict):
+    masked_text = input_text
+    for entity in sorted(entities, key=lambda x: x['start'], reverse=True):
+        if entity['entity_group'] not in ['CARDINAL', 'EVENT']:
+            masked_text = masked_text[:entity['start']] + mask_dict[entity['word']] + masked_text[entity['end']:]
+    return masked_text
+
 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)
@@ -157,7 +163,7 @@ if Run_Button and input_text:
     # Create mask dictionary
     mask_dict = create_mask_dict(output_comb)
 
-    masked_text = " ".join(entity['masked_word'] for entity in output_comb)
+    masked_text = create_masked_text(input_text, output_comb, mask_dict)
     
     # Apply masking and add masked_word column
     for entity in output_comb:
@@ -183,20 +189,36 @@ if Run_Button and input_text:
         else:
             label = entity['entity_group']
         spacy_display["ents"].append({"start": entity["start"], "end": entity["end"], "label": label})
+        # Custom CSS to prevent label overlap
+    custom_css = """
+    <style>
+    .entity-label {
+        font-size: 0.7em;
+        line-height: 1;
+        padding: 0.25em;
+        border-radius: 0.25em;
+        top: -1.5em;
+        position: relative;
+    }
+    </style>
+    """
     
-    html = spacy.displacy.render(spacy_display, style="ent", minify=True, manual=True)
+    html = custom_css + spacy.displacy.render(spacy_display, style="ent", minify=True, manual=True)
     st.write(html, unsafe_allow_html=True)
 
+ # Download button
     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"
-            )
+    masked_file_content = export_masked_text(masked_text, export_file_type)
+    if masked_file_content:
+        st.download_button(
+            label="Download Masked Text",
+            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)
+    st.json(mask_dict)
+
+    st.subheader("Masked Text Preview")
+    st.text(masked_text)