testin the trick

modeling_stacked.py

@@ -53,7 +53,7 @@ class ExtendedMultitaskModelForTokenClassification(PreTrainedModel):
 
         # Load floret model
         self.dummy_param = nn.Parameter(torch.zeros(1))
-        self.safe_floret = SafeFloretWrapper(self.config.config.filename)
+        self.model_floret = floret.load_model(self.config.config.filename)
         # print(self.config.config)
         print(type(self.config))
         print(self.config.config.filename)
@@ -66,7 +66,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}, ---, {type(input_ids)} ----- {type(self.safe_floret)}"
+            f"Check if it arrives here: {input_ids}, ---, {type(input_ids)} ----- {type(self.model_floret)}"
         )
         if isinstance(input_ids, str):
             # If the input is a single string, make it a list for floret
@@ -77,7 +77,7 @@ class ExtendedMultitaskModelForTokenClassification(PreTrainedModel):
             raise ValueError(f"Unexpected input type: {type(input_ids)}")
 
         # Use the SafeFloretWrapper to get predictions
-        predictions, probabilities = self.safe_floret.predict(texts)
+        predictions, probabilities = self.model_floret.predict(texts, k=1)
         print(f"Predictions: {predictions}")
         print(f"Probabilities: {probabilities}")
         # print(self.model_floret(input_ids))
@@ -94,7 +94,7 @@ class ExtendedMultitaskModelForTokenClassification(PreTrainedModel):
         #     return torch.tensor(predictions)
         # else:
         # If no text is found, return dummy output
-        return torch.zeros((1, 2))  # Dummy tensor with shape (batch_size, num_classes)
+        return predictions  # Dummy tensor with shape (batch_size, num_classes)
 
     def state_dict(self, *args, **kwargs):
         # Return an empty state dictionary