Update modeling_st2.py

modeling_st2.py

@@ -33,9 +33,9 @@ class ST2ModelV2(nn.Module):
         super(ST2ModelV2, self).__init__()
         self.args = args
 
-        self.config = AutoConfig.from_pretrained(args.model_name_or_path)
-        self.model = AutoModel.from_pretrained(args.model_name_or_path)
-        self.tokenizer = AutoTokenizer.from_pretrained(args.model_name_or_path)
+        self.config = AutoConfig.from_pretrained("roberta-large")
+        self.model = AutoModel.from_pretrained("roberta-large")
+        self.tokenizer = AutoTokenizer.from_pretrained("roberta-large")
 
         classifier_dropout = self.args.dropout
         self.dropout = nn.Dropout(classifier_dropout)
@@ -104,40 +104,18 @@ class ST2ModelV2(nn.Module):
                 return_dict=return_dict,
             )
         else:
-            if self.args.model_name_or_path in ['facebook/bart-large', 'facebook/bart-base', 'facebook/bart-large-cnn']:
-                outputs = self.model(
-                    input_ids,
-                    attention_mask=attention_mask,
-                    head_mask=head_mask,
-                    inputs_embeds=inputs_embeds,
-                    output_attentions=output_attentions,
-                    output_hidden_states=output_hidden_states,
-                    return_dict=return_dict,
-                )
-            elif self.args.model_name_or_path in ['microsoft/deberta-base']:
-                outputs = self.model(
-                    input_ids,
-                    attention_mask=attention_mask,
-                    token_type_ids=token_type_ids,
-                    position_ids=position_ids,
-                    inputs_embeds=inputs_embeds,
-                    output_attentions=output_attentions,
-                    output_hidden_states=output_hidden_states,
-                    return_dict=return_dict,
-                )
-
-            else:               
-                outputs = self.model(
-                    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,
-                )
+            
+            outputs = self.model(
+                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,
+            )
 
         sequence_output = outputs[0]
         if signal_bias_mask is not None and self.args.signal_bias_on_top_of_lm: