degbu

generic_ner.py

@@ -256,6 +256,22 @@ def postprocess_entities(entities):
     return filtered_entities
 
 
+def remove_included_entities(entities):
+    # Loop through entities and remove those whose text is included in another with the same label
+    final_entities = []
+    for i, entity in enumerate(entities):
+        is_included = False
+        for other_entity in entities:
+            if entity != other_entity and entity["entity"] == other_entity["entity"]:
+                # Check if entity's text is a substring of another entity's text
+                if entity["text"] in other_entity["text"]:
+                    is_included = True
+                    break
+        if not is_included:
+            final_entities.append(entity)
+    return final_entities
+
+
 class MultitaskTokenClassificationPipeline(Pipeline):
 
     def _sanitize_parameters(self, **kwargs):
@@ -381,10 +397,14 @@ class MultitaskTokenClassificationPipeline(Pipeline):
         )
 
         print("After 1:")
+        all_entities = remove_included_entities(all_entities)
+        pprint(all_entities)
+
+        print("After 2:")
         pprint(all_entities)
         # Attach "comp.function" entities to the closest non-"comp.function" entity
         all_entities = attach_comp_to_closest(all_entities)
-        print("After 2:")
+        print("After 3:")
         pprint(all_entities)
         print("\n")
         return all_entities