dd

generic_ner.py

@@ -687,6 +687,7 @@ def remove_trailing_stopwords(entities):
         print(f"Remained entities: {len(new_entities)}")
     return new_entities
 
+
 class MultitaskTokenClassificationPipeline(Pipeline):
 
     def _sanitize_parameters(self, **kwargs):
@@ -717,6 +718,7 @@ class MultitaskTokenClassificationPipeline(Pipeline):
         attention_mask = torch.tensor([inputs["attention_mask"]], dtype=torch.long).to(
             self.model.device
         )
+        print(f"Let's check the model: {self.model}")
         with torch.no_grad():
             outputs = self.model(input_ids, attention_mask)
         return outputs, text_sentences, text
@@ -755,12 +757,6 @@ class MultitaskTokenClassificationPipeline(Pipeline):
 
             entities[task] = get_entities(words_list, preds_list, confidence_list, text)
 
-        # add titles to comp entities
-        # from pprint import pprint
-
-        # print("Before:")
-        # pprint(entities)
-
         all_entities = []
         coarse_entities = []
         for key in entities: