degbu

generic_ner.py

@@ -274,18 +274,15 @@ class MultitaskTokenClassificationPipeline(Pipeline):
         """Check if entity1 is fully within the bounds of entity2."""
         return entity1["start"] >= entity2["start"] and entity1["end"] <= entity2["end"]
 
-    def postprocess_entities(self, ner_results, text_sentence):
+    def postprocess_entities(self, all_entities, text_sentence):
         # Collect all entities in one list for processing
-        all_entities = []
-        for key in ner_results:
-            if key not in ["NE-COARSE-LIT"]:
-                all_entities.extend(ner_results[key])
-                # for entity in ner_results[key]:
-                #     print(
-                #         text_sentence[entity["start"] : entity["end"]],
-                #         "--",
-                #         entity["word"],
-                #     )
+
+        # for entity in ner_results[key]:
+        #     print(
+        #         text_sentence[entity["start"] : entity["end"]],
+        #         "--",
+        #         entity["word"],
+        #     )
 
         # Sort entities by start position, then by end position (to handle nested structures)
         all_entities.sort(key=lambda x: (x["start"], -x["end"]))
@@ -360,10 +357,18 @@ class MultitaskTokenClassificationPipeline(Pipeline):
 
         print("Before:")
         pprint(entities)
-        postprocessed_entities = self.postprocess_entities(entities, text_sentence)
 
+        all_entities = []
+        for key in entities:
+            if key not in ["NE-COARSE-LIT"]:
+                all_entities.extend(entities[key])
+
+        postprocessed_entities = self.postprocess_entities(all_entities, text_sentence)
+        print("After 1:")
+        pprint(postprocessed_entities)
         # Attach "comp.function" entities to the closest non-"comp.function" entity
         postprocessed_entities = attach_comp_to_closest(postprocessed_entities)
-        print("After:")
+        print("After 2:")
         pprint(postprocessed_entities)
+        print("\n")
         return postprocessed_entities