Update app.py

app.py

@@ -93,18 +93,20 @@ def entity_comb(output):
             output_comb.append(entity)
     return output_comb
 
-def create_mask_dict(entities):
-    mask_dict = {}
-    entity_counters = {}
-    for entity in entities:
+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']:
-            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
+            # 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
+                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):