dd

generic_ner.py

@@ -718,7 +718,9 @@ 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}")
+        # print(f"Let's check the model: {self.model}")
+        # check get floret model
+        print(f"Let's check the model: {self.model.get_floret_model()}")
         with torch.no_grad():
             outputs = self.model(input_ids, attention_mask)
         return outputs, text_sentences, text

modeling_stacked.py

@@ -60,6 +60,9 @@ class ExtendedMultitaskModelForTokenClassification(PreTrainedModel):
         # Initialize weights and apply final processing
         self.post_init()
 
+    def get_floret_model(self):
+        return self.model_floret
+
     def forward(
         self,
         input_ids: Optional[torch.Tensor] = None,