Update app.py

app.py

@@ -1,6 +1,5 @@
 import streamlit as st
 import pandas as pd
-import spacy
 from transformers import pipeline, AutoModelForTokenClassification, AutoTokenizer
 import PyPDF2
 import docx
@@ -92,6 +91,7 @@ def entity_comb(output):
         else:
             output_comb.append(entity)
     return output_comb
+
 def create_mask_dict(entities):
     mask_dict = {}
     entity_counters = {}
@@ -104,14 +104,13 @@ 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):
-    # Create the mask dictionary
     mask_dict = create_mask_dict(entities)
     
     masked_text = input_text
     for entity in sorted(entities, key=lambda x: x['start'], reverse=True):
         if entity['entity_group'] not in ['CARDINAL', 'EVENT']:
-            # Replace the entity with its entity group from the mask dictionary
             masked_text = (
                 masked_text[:entity['start']] +
                 f"<{mask_dict[entity['word']]}> " +  # Use angle brackets for clarity
@@ -140,47 +139,17 @@ if Run_Button and input_text:
                 entity['end'] += offset
         
         all_outputs.extend(output)
-        
     
     # Combine entities
-    
     output_comb = entity_comb(all_outputs)
     
-    # Create mask dictionary
-    mask_dict = create_mask_dict(output_comb)
-
+    # Create masked text and masking dictionary
     masked_text = create_masked_text(input_text, output_comb)
-    
-    # Apply masking and add masked_word column
-    for entity in output_comb:
-        if entity['entity_group'] not in ['CARDINAL', 'EVENT']:
-            entity['masked_word'] = mask_dict.get(entity['word'], entity['word'])
-        else:
-            entity['masked_word'] = entity['word']
-    print("output_comb", output_comb)
-    #df = pd.DataFrame.from_dict(output_comb)
-    #cols_to_keep = ['word', 'entity_group', 'score', 'start', 'end']
-    #df_final = df[cols_to_keep].loc[:,~df.columns.duplicated()].copy()
-    
-    #st.subheader("Recognized Entities")
-    #st.dataframe(df_final)
+    mask_dict = create_mask_dict(output_comb)
 
-    
-    
-    # Spacy display logic with entity numbering
-    spacy_display = {"ents": [], "text": input_text, "title": None}
-    for entity in output_comb:
-        if entity['entity_group'] not in ['CARDINAL', 'EVENT']:
-            label = f"{entity['entity_group']}_{mask_dict[entity['word']].split('_')[1]}"
-        else:
-            label = entity['entity_group']
-        spacy_display["ents"].append({"start": entity["start"], "end": entity["end"], "label": label})
-    
-    html = spacy.displacy.render(spacy_display, style="ent", minify=True, manual=True)
-    st.write(html, unsafe_allow_html=True)
+    # Display the masked text and masking dictionary
+    st.subheader("Masked Text Preview")
+    st.text(masked_text)
 
     st.subheader("Masking Dictionary")
     st.json(mask_dict)
-
-    st.subheader("Masked Text Preview")
-    st.text(masked_text)