dd

generic_ner.py

@@ -712,19 +712,8 @@ class MultitaskTokenClassificationPipeline(Pipeline):
 
     def _forward(self, inputs):
         inputs, text_sentences, text = inputs
-        input_ids = torch.tensor([inputs["input_ids"]], dtype=torch.long).to(
-            self.model.device
-        )
-        attention_mask = torch.tensor([inputs["attention_mask"]], dtype=torch.long).to(
-            self.model.device
-        )
-        # print(f"Let's check the model: {self.model}")
-        # check get floret model
-
-        with torch.no_grad():
-            outputs = self.model(input_ids, attention_mask)
 
-        return outputs, text_sentences, text
+        return text
 
     def is_within(self, entity1, entity2):
         """Check if entity1 is fully within the bounds of entity2."""
@@ -733,58 +722,13 @@ class MultitaskTokenClassificationPipeline(Pipeline):
             and entity1["rOffset"] <= entity2["rOffset"]
         )
 
-    def postprocess(self, outputs, **kwargs):
+    def postprocess(self, text, **kwargs):
         """
         Postprocess the outputs of the model
         :param outputs:
         :param kwargs:
         :return:
         """
-        tokens_result, text_sentence, text = outputs
-
-        predictions = {}
-        confidence_scores = {}
-        for task, logits in tokens_result.logits.items():
-            predictions[task] = torch.argmax(logits, dim=-1).tolist()[0]
-            confidence_scores[task] = F.softmax(logits, dim=-1).tolist()[0]
-
-        entities = {}
-        for task in predictions.keys():
-            words_list, preds_list, confidence_list = realign(
-                text_sentence,
-                predictions[task],
-                confidence_scores[task],
-                self.tokenizer,
-                self.id2label[task],
-            )
-
-            entities[task] = get_entities(words_list, preds_list, confidence_list, text)
-
-        all_entities = []
-        coarse_entities = []
-        for key in entities:
-            if key in ["NE-COARSE-LIT"]:
-                coarse_entities = entities[key]
-            all_entities.extend(entities[key])
-
-        if DEBUG:
-            print(all_entities)
-        # print("After remove_included_entities:")
-        all_entities = remove_included_entities(all_entities)
-        if DEBUG:
-            print("After remove_included_entities:", all_entities)
-        all_entities = remove_trailing_stopwords(all_entities)
-        if DEBUG:
-            print("After remove_trailing_stopwords:", all_entities)
-        all_entities = postprocess_entities(all_entities)
-        if DEBUG:
-            print("After postprocess_entities:", all_entities)
-        all_entities = refine_entities_with_coarse(all_entities, coarse_entities)
-        if DEBUG:
-            print("After refine_entities_with_coarse:", all_entities)
-        # print("After attach_comp_to_closest:")
-        # pprint(all_entities)
-        # print("\n")
 
         print(f"Let's check the model: {self.model.get_floret_model()}")
         predictions, probabilities = self.model.get_floret_model().predict([text], k=1)