Update app.py

app.py

@@ -105,11 +105,11 @@ 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, mask_dict):
+def create_masked_text(input_text, entities, mask_dict):
     masked_text = input_text
-    for word, mask in sorted(mask_dict.items(), key=lambda x: len(x[0]), reverse=True):
-        masked_text = re.sub(r'\b' + re.escape(word) + r'\b', mask, masked_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