Update XLMRoBERTaClassifier.py

XLMRoBERTaClassifier.py

@@ -15,7 +15,7 @@ warnings.filterwarnings("ignore")
 
 class XLMRoBERTaClassifier(PreTrainedModel):
   def __init__(self, dropout=0.3, model_name='xlm-roberta-large'):
-    self.config = AutoConfig.from_pretrained("hugo-albert/xlm-robertargument")
+    self.config = AutoConfig.from_pretrained("FacebookAI/xlm-roberta-large")
     super(XLMRoBERTaClassifier, self).__init__(self.config)
     self.roberta = XLMRobertaModel.from_pretrained(model_name)
     self.dropout = nn.Dropout(dropout)
@@ -29,9 +29,8 @@ class XLMRoBERTaClassifier(PreTrainedModel):
     self.final_layer = nn.Linear(128, 1)
 
   def forward(self, input_ids, attention_mask): #, extra_features):
-    input_id, mask = torch.Tensor(input_ids).long(), torch.LongTensor(attention_mask).long()
-    roberta_output = self.roberta(input_ids = input_id,
-    attention_mask=mask)
+    roberta_output = self.roberta(input_ids = input_ids,
+    attention_mask=attention_mask)
     last_hidden_state = roberta_output.last_hidden_state
     conv_output = self.conv1(last_hidden_state)
     pool_output = self.pool(conv_output)
@@ -44,4 +43,4 @@ class XLMRoBERTaClassifier(PreTrainedModel):
     final_output = self.final_layer(dropout_output)
     sigmoid_output = self.sigmoid(final_output)
     sigmoid_output = torch.squeeze(sigmoid_output)
-    return {"logits": sigmoid_output}
+    return sigmoid_output