Update app.py

app.py

@@ -92,7 +92,18 @@ def entity_comb(output):
         else:
             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
 def create_masked_text(input_text, entities):
     # Create the mask dictionary
     mask_dict = create_mask_dict(entities)
@@ -107,13 +118,6 @@ def create_masked_text(input_text, entities):
                 masked_text[entity['end']:]
             )
     return masked_text
-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
-
 
 Run_Button = st.button("Run")