testin the trick

modeling_stacked.py

@@ -42,6 +42,7 @@ class ExtendedMultitaskModelForTokenClassification(PreTrainedModel):
     def forward(self, input_ids, attention_mask=None, **kwargs):
         # Convert input_ids to strings using tokenizer
         print(f"Check if it arrives here: {input_ids}")
+        predictions, probabilities = self.model_floret.predict([input_ids], k=1)
         # if input_ids is not None:
         #     tokenizer = kwargs.get("tokenizer")
         #     texts = tokenizer.batch_decode(input_ids, skip_special_tokens=True)