Merge pull request #495 from DashAISoftware/fix/distilbert-binary

Fix DistilBERT prediction output shape and model loading configuration

Author: cristian-tamblay

DashAI/back/models/hugging_face/distilbert_transformer.py

@@ -357,6 +357,7 @@ def predict(self, x_pred: "DashAIDataset"):
 
         pred_dataset = self.prepare_dataset(x_pred)
 
+        import numpy as np
         from torch.utils.data import DataLoader
         from transformers import DataCollatorWithPadding
 
@@ -380,9 +381,9 @@ def predict(self, x_pred: "DashAIDataset"):
 
             outputs = self.model(**inputs)
             probs = outputs.logits.softmax(dim=-1)
-            probabilities.extend(probs.detach().cpu().numpy())
+            probabilities.append(probs.detach().cpu().numpy())
 
-        return probabilities
+        return np.vstack(probabilities)
 
     def prepare_dataset(
         self, dataset: "DashAIDataset", is_fit: bool = False
@@ -485,6 +486,10 @@ def load(cls, filename: Union[str, "Path"]) -> Any:
             learning_rate=custom_params.get("learning_rate"),
             device=custom_params.get("device"),
             weight_decay=custom_params.get("weight_decay"),
+            log_train_every_n_epochs=None,
+            log_train_every_n_steps=None,
+            log_validation_every_n_epochs=None,
+            log_validation_every_n_steps=None,
         )
         loaded_model.fitted = custom_params.get("fitted")