Update modeling_xlm_roberta_for_glue.py

modeling_xlm_roberta_for_glue.py

@@ -6,16 +6,16 @@ from torch.nn import CrossEntropyLoss, MSELoss, BCEWithLogitsLoss
 from transformers.modeling_outputs import SequenceClassifierOutput, QuestionAnsweringModelOutput, TokenClassifierOutput
 
 from .modeling_bert import XLMRobertaPreTrainedModel, XLMRobertaModel
-from .configuration_bert import JinaBertConfig
+from .configuration_xlm_roberta import XLMRobertaFlashConfig
 
 
 class XLMRobertaForSequenceClassification(XLMRobertaPreTrainedModel):
-    def __init__(self, config: JinaBertConfig):
+    def __init__(self, config: XLMRobertaFlashConfig):
         super().__init__(config)
         self.num_labels = config.num_labels
         self.config = config
 
-        self.bert = XLMRobertaModel(config)
+        self.roberta = XLMRobertaModel(config)
         classifier_dropout = (
             config.classifier_dropout
             if config.classifier_dropout is not None
@@ -56,11 +56,16 @@ class XLMRobertaForSequenceClassification(XLMRobertaPreTrainedModel):
         assert output_attentions is None
         assert output_hidden_states is None
         assert return_dict
-        outputs = self.bert(
+        outputs = self.roberta(
             input_ids,
             attention_mask=attention_mask,
             token_type_ids=token_type_ids,
             position_ids=position_ids,
+            head_mask=head_mask,
+            inputs_embeds=inputs_embeds,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
         )
 
         pooled_output = outputs[1]