umarigan commited on
Commit
5f73d00
·
verified ·
1 Parent(s): 861ad26

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +13 -11
app.py CHANGED
@@ -93,18 +93,20 @@ def entity_comb(output):
93
  output_comb.append(entity)
94
  return output_comb
95
 
96
- def create_mask_dict(entities):
97
- mask_dict = {}
98
- entity_counters = {}
99
- for entity in entities:
 
 
100
  if entity['entity_group'] not in ['CARDINAL', 'EVENT']:
101
- if entity['word'] not in mask_dict:
102
- if entity['entity_group'] not in entity_counters:
103
- entity_counters[entity['entity_group']] = 1
104
- else:
105
- entity_counters[entity['entity_group']] += 1
106
- mask_dict[entity['word']] = f"{entity['entity_group']}_{entity_counters[entity['entity_group']]}"
107
- return mask_dict
108
  def create_masked_text(input_text, entities, mask_dict):
109
  masked_text = input_text
110
  for entity in sorted(entities, key=lambda x: x['start'], reverse=True):
 
93
  output_comb.append(entity)
94
  return output_comb
95
 
96
+ def create_masked_text(input_text, entities):
97
+ # Create the mask dictionary
98
+ mask_dict = create_mask_dict(entities)
99
+
100
+ masked_text = input_text
101
+ for entity in sorted(entities, key=lambda x: x['start'], reverse=True):
102
  if entity['entity_group'] not in ['CARDINAL', 'EVENT']:
103
+ # Replace the entity with its entity group from the mask dictionary
104
+ masked_text = (
105
+ masked_text[:entity['start']] +
106
+ f"<{mask_dict[entity['word']]}> " + # Use angle brackets for clarity
107
+ masked_text[entity['end']:]
108
+ )
109
+ return masked_text
110
  def create_masked_text(input_text, entities, mask_dict):
111
  masked_text = input_text
112
  for entity in sorted(entities, key=lambda x: x['start'], reverse=True):