Make missing lables handling more direct

Author: MGlauer

chebai/models/electra.py

@@ -299,8 +299,7 @@ def _process_for_loss(
             labels = labels.float()
         if "missing_labels" in kwargs_copy:
             missing_labels = kwargs_copy.pop("missing_labels")
-            output = output * (~missing_labels).int() - 10000 * missing_labels.int()
-            labels = labels * (~missing_labels).int()
+            output[missing_labels] = -1e8
         if self.model_type == "classification":
             assert ((labels <= torch.tensor(1.0)) & (labels >= torch.tensor(0.0))).all()
         return output, labels, kwargs_copy
@@ -332,7 +331,7 @@ def _get_prediction_and_labels(
             if "missing_labels" in loss_kwargs:
                 #print('bla')
                 missing_labels = loss_kwargs["missing_labels"]
-                d = d * (~missing_labels).int().to(device=d.device)
+                d[missing_labels] = 0
             return d, labels.int() if labels is not None else None
         elif self.model_type == 'regression':
             return d, labels

configs/metrics/micro-macro-f1-roc-auc.yml

@@ -5,6 +5,7 @@ init_args:
       class_path: torchmetrics.classification.MultilabelF1Score
       init_args:
         average: micro
+        zero_division: 1
     macro-f1:
       class_path: chebai.callbacks.epoch_metrics.MacroF1
     roc-auc:

configs/model/electra.yml

@@ -1,6 +1,5 @@
 class_path: chebai.models.Electra
 init_args:
-  model_type: regression
   optimizer_kwargs:
     lr: 1e-4
   config: