add comp back

generic_ner.py

@@ -500,20 +500,25 @@ def refine_entities_with_coarse(all_entities, coarse_entities):
 
     # 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:
             coarse_entity = coarse_lookup[key]
             # 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
-
+                entity["type"] = coarse_entity[
+                    "type"
+                ]  # Update the type if the confidence is higher
+            else:
+                entity["type"] = entity["type"].split('.')[0]
     # No need to append to refined_entities, we're modifying in place
     return all_entities
 
 
-
 def remove_trailing_stopwords(entities):
     """
     This function removes stopwords and punctuation from both the beginning and end of each entity's text