add comp back

generic_ner.py

@@ -486,6 +486,45 @@ def remove_included_entities(entities):
     return final_entities
 
 
+def refine_entities_with_coarse(all_entities, coarse_entities):
+    """
+    Looks through all entities and refines them based on the coarse entities.
+    If a surface match is found in the coarse entities and the types match,
+    the entity with the higher confidence_ner is kept.
+    """
+    refined_entities = []
+
+    # Create a dictionary for coarse entities based on surface and type for quick lookup
+    coarse_lookup = {}
+    for coarse_entity in coarse_entities:
+        key = (coarse_entity["surface"], coarse_entity["type"])
+        coarse_lookup[key] = coarse_entity
+
+    # Iterate through all entities and compare with the coarse entities
+    for entity in all_entities:
+        key = (
+            entity["surface"],
+            entity["type"].split(".")[0],
+        )  # Use the coarse type for comparison
+
+        if key in coarse_lookup:
+            # If the types match, compare confidence_ner and keep the one with the higher confidence
+            coarse_entity = coarse_lookup[key]
+            if entity["confidence_ner"] > coarse_entity["confidence_ner"]:
+                refined_entities.append(
+                    entity
+                )  # Keep the current entity with higher confidence
+            else:
+                refined_entities.append(
+                    coarse_entity
+                )  # Keep the coarse entity with higher confidence
+        else:
+            # If no match in coarse, just add the entity to refined entities
+            refined_entities.append(entity)
+
+    return refined_entities
+
+
 def remove_trailing_stopwords(entities):
     """
     This function removes stopwords and punctuation from both the beginning and end of each entity's text
@@ -715,8 +754,10 @@ class MultitaskTokenClassificationPipeline(Pipeline):
         # pprint(entities)
 
         all_entities = []
+        coarse_entities = []
         for key in entities:
-            # if key not in ["NE-COARSE-LIT"]:
+            if key in ["NE-COARSE-LIT"]:
+                coarse_entities = entities[key]
             all_entities.extend(entities[key])
 
         if DEBUG:
@@ -725,7 +766,7 @@ class MultitaskTokenClassificationPipeline(Pipeline):
         all_entities = remove_included_entities(all_entities)
         all_entities = remove_trailing_stopwords(all_entities)
         all_entities = postprocess_entities(all_entities)
-
+        all_entities = refine_entities_with_coarse(all_entities, coarse_entities)
         # print("After attach_comp_to_closest:")
         # pprint(all_entities)
         # print("\n")