Update app.py

app.py

@@ -93,6 +93,19 @@ def entity_comb(output):
             output_comb.append(entity)
     return output_comb
 
+def create_mask_dict(entities):
+    mask_dict = {}
+    entity_counters = {}
+    for entity in entities:
+        if entity['entity_group'] not in ['CARDINAL', 'EVENT']:
+            if entity['word'] not in mask_dict:
+                if entity['entity_group'] not in entity_counters:
+                    entity_counters[entity['entity_group']] = 1
+                else:
+                    entity_counters[entity['entity_group']] += 1
+                mask_dict[entity['word']] = f"{entity['entity_group']}_{entity_counters[entity['entity_group']]}"
+    return mask_dict
+
 Run_Button = st.button("Run")
 
 if Run_Button and input_text:
@@ -118,17 +131,34 @@ if Run_Button and input_text:
     # Combine entities
     output_comb = entity_comb(all_outputs)
     
+    # Create mask dictionary
+    mask_dict = create_mask_dict(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']
+    
     df = pd.DataFrame.from_dict(output_comb)
-    cols_to_keep = ['word', 'entity_group', 'score', 'start', 'end']
+    cols_to_keep = ['word', 'masked_word', 'entity_group', 'score', 'start', 'end']
     df_final = df[cols_to_keep]
     
     st.subheader("Recognized Entities")
     st.dataframe(df_final)
     
-    # Spacy display logic
+    # Spacy display logic with entity numbering
     spacy_display = {"ents": [], "text": input_text, "title": None}
     for entity in output_comb:
-        spacy_display["ents"].append({"start": entity["start"], "end": entity["end"], "label": entity["entity_group"]})
+        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)
+    st.write(html, unsafe_allow_html=True)
+
+    st.subheader("Masking Dictionary")
+    st.json(mask_dict)