Fix electra pretraining

Author: MGlauer

chebai/loss/pretraining.py

@@ -0,0 +1,17 @@
+import torch
+
+class ElectraPreLoss(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.ce = torch.nn.CrossEntropyLoss()
+
+    def forward(self, input, target, **loss_kwargs):
+        t, p = input
+        gen_pred, disc_pred = t
+        gen_tar, disc_tar = p
+        gen_loss = self.ce(target=torch.argmax(gen_tar.int(), dim=-1), input=gen_pred)
+        disc_loss = self.ce(
+            target=torch.argmax(disc_tar.int(), dim=-1), input=disc_pred
+        )
+        return gen_loss + disc_loss
+

chebai/models/base.py

@@ -41,10 +41,13 @@ def __init_subclass__(cls, **kwargs):
     def _get_prediction_and_labels(self, data, labels, output):
         return output, labels
 
+    def _process_labels_in_batch(self, batch):
+        return batch.y.float()
+
     def _process_batch(self, batch, batch_idx):
         return dict(
             features=batch.x,
-            labels=batch.y.float(),
+            labels=self._process_labels_in_batch(batch),
             model_kwargs=batch.additional_fields["model_kwargs"],
             loss_kwargs=batch.additional_fields["loss_kwargs"],
             idents=batch.additional_fields["idents"],

chebai/models/electra.py

@@ -47,15 +47,15 @@ def __init__(self, config=None, **kwargs):
     def as_pretrained(self):
         return self.discriminator
 
-    def _process_batch(self, batch, batch_idx):
-        return dict(features=batch.x, labels=None, mask=batch.mask)
+    def _process_labels_in_batch(self, batch):
+        return None
 
-    def forward(self, data):
+    def forward(self, data, **kwargs):
         features = data["features"]
         self.batch_size = batch_size = features.shape[0]
         max_seq_len = features.shape[1]
 
-        mask = data["mask"]
+        mask = kwargs["mask"]
         with torch.no_grad():
             dis_tar = (
                 torch.rand((batch_size,), device=self.device) * torch.sum(mask, dim=-1)
@@ -96,25 +96,6 @@ def forward(self, data):
     def _get_prediction_and_labels(self, batch, labels, output):
         return torch.softmax(output[0][1], dim=-1), output[1][1].int()
 
-    def _get_data_for_loss(self, model_output, labels):
-        return dict(input=model_output, target=None)
-
-
-class ElectraPreLoss(torch.nn.Module):
-    def __init__(self):
-        super().__init__()
-        self.ce = torch.nn.CrossEntropyLoss()
-
-    def forward(self, input, target):
-        t, p = input
-        gen_pred, disc_pred = t
-        gen_tar, disc_tar = p
-        gen_loss = self.ce(target=torch.argmax(gen_tar.int(), dim=-1), input=gen_pred)
-        disc_loss = self.ce(
-            target=torch.argmax(disc_tar.int(), dim=-1), input=disc_pred
-        )
-        return gen_loss + disc_loss
-
 
 def filter_dict(d, filter_key):
     return {

configs/model/electra-for-pretraining.yml

@@ -0,0 +1,20 @@
+class_path: chebai.models.ElectraPre
+init_args:
+  criterion:
+    class_path: chebai.loss.pretraining.ElectraPreLoss
+  out_dim: null
+  optimizer_kwargs:
+    lr: 1e-4
+  config:
+    generator:
+      vocab_size: 1400
+      max_position_embeddings: 1800
+      num_attention_heads: 8
+      num_hidden_layers: 6
+      type_vocab_size: 1
+    discriminator:
+      vocab_size: 1400
+      max_position_embeddings: 1800
+      num_attention_heads: 8
+      num_hidden_layers: 6
+      type_vocab_size: 1