Updated ganbert.py

ganbert.py

@@ -76,10 +76,24 @@ class GAN(PreTrainedModel):
         self.generator.cuda()
         self.discriminator.cuda()
         self.transformer.cuda()
-  def forward(self,b_input_ids,b_input_mask):
+  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,
+    ):
     # Encode real data in the Transformer
-    real_batch_size = b_input_ids.shape[0]
-    model_outputs = self.transformer(b_input_ids, attention_mask=b_input_mask)
+    real_batch_size = input_ids.shape[0]
+    model_outputs = self.transformer(input_ids, attention_mask=attention_mask)
     # print('got transformer output')
     
     hidden_states = torch.mean(model_outputs[0],dim=1)
@@ -88,7 +102,3 @@ class GAN(PreTrainedModel):
     disciminator_input = torch.cat([hidden_states, gen_rep], dim=0)
     features, logits, probs = self.discriminator(disciminator_input)
     return model_outputs[0]
-
-if __name__ == '__main__':
-  ganconfig = GanBertConfig()
-  clickbaitmodel = GAN(ganconfig)