add comp back

generic_ner.py

@@ -490,39 +490,28 @@ 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.
+    the entity's confidence_ner and type are updated based on the coarse entity.
     """
-    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"])
+        key = (coarse_entity["surface"], coarse_entity["type"].split(".")[0])
         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
+        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)
+            # If a match is found, update the confidence_ner and type in the entity
+            if entity["confidence_ner"] < coarse_entity["confidence_ner"]:
+                entity["confidence_ner"] = coarse_entity["confidence_ner"]
+                entity["type"] = coarse_entity["type"]  # Update the type if the confidence is higher
+
+    # No need to append to refined_entities, we're modifying in place
+    return all_entities
 
-    return refined_entities
 
 
 def remove_trailing_stopwords(entities):