Updated ganbert.py

ganbert.py

@@ -30,6 +30,7 @@ from transformers import (
     set_seed,
     get_constant_schedule_with_warmup,
     Trainer,TrainingArguments,EarlyStoppingCallback)
+
 from datasets import Dataset
 import torch.nn as nn
 import torch.nn.functional as F
@@ -77,29 +78,14 @@ class GAN(PreTrainedModel):
         self.generator.cuda()
         self.discriminator.cuda()
         self.transformer.cuda()
-  def forward(self,
-        input_ids: Optional[torch.Tensor] = None,
-        attention_mask: Optional[torch.Tensor] = None,
-        token_type_ids: Optional[torch.Tensor] = None,
-        position_ids: Optional[torch.Tensor] = None,
-        head_mask: Optional[torch.Tensor] = None,
-        inputs_embeds: Optional[torch.Tensor] = None,
-        encoder_hidden_states: Optional[torch.Tensor] = None,
-        encoder_attention_mask: Optional[torch.Tensor] = None,
-        past_key_values: Optional[List[torch.FloatTensor]] = None,
-        use_cache: Optional[bool] = None,
-        output_attentions: Optional[bool] = None,
-        output_hidden_states: Optional[bool] = None,
-        return_dict: Optional[bool] = None,
-    ):
+  def forward(self,**kwargs):
     # Encode real data in the Transformer
-    real_batch_size = input_ids.shape[0]
-    model_outputs = self.transformer(input_ids, attention_mask=attention_mask)
+    # real_batch_size = input_ids.shape[0]
+    model_outputs = self.transformer(**kwargs)
     # print('got transformer output')
-    
-    hidden_states = torch.mean(model_outputs[0],dim=1)
-    noise = torch.zeros(real_batch_size, self.ns, device=self.dv).uniform_(0, 1).to(self.dv)
-    gen_rep = self.generator(noise)
-    disciminator_input = torch.cat([hidden_states, gen_rep], dim=0)
-    features, logits, probs = self.discriminator(disciminator_input)
-    return model_outputs[0]
+    # hidden_states = torch.mean(model_outputs[0],dim=1)
+    # noise = torch.zeros(real_batch_size, self.ns, device=self.dv).uniform_(0, 1).to(self.dv)
+    # gen_rep = self.generator(noise)
+    # disciminator_input = torch.cat([hidden_states, gen_rep], dim=0)
+    # features, logits, probs = self.discriminator(disciminator_input)
+    return model_outputs