Upload AbLang

model.py

@@ -35,13 +35,23 @@ class AbLang(PreTrainedModel):
         self.AbEmbeddings = AbEmbeddings(config)    
         self.EncoderBlocks = EncoderBlocks(config)
         
-    def forward(self, inputs):
-        src = self.AbEmbeddings(inputs['input_ids'])
-        outputs = self.EncoderBlocks(src, attention_mask=1-inputs['attention_mask'], output_attentions=False)
-        return apply_cls_embeddings(inputs, outputs)
+    def forward(
+            self, 
+            input_ids=None,
+            attention_mask=None,
+            token_type_ids=None,
+            output_attentions=None,
+            output_hidden_states=None,
+        ):
+        src = self.AbEmbeddings(input_ids)
+        outputs = self.EncoderBlocks(src, 
+                                     attention_mask=1-attention_mask, 
+                                     output_attentions=output_attentions, 
+                                     output_hidden_states=output_hidden_states)
+        return apply_cls_embeddings(attention_mask, outputs)
     
-def apply_cls_embeddings(inputs, outputs):
-    mask = inputs['attention_mask'].float()
+def apply_cls_embeddings(attention_mask, outputs):
+    mask = attention_mask.float()
     d = {k: v for k, v in torch.nonzero(mask).cpu().numpy()} # dict of sep tokens
     # make sep token invisible
     for i in d: