Update custom_st.py

custom_st.py

@@ -51,8 +51,8 @@ class Transformer(nn.Module):
         if config_args is None:
             config_args = {}
 
+        
         self.config = AutoConfig.from_pretrained(model_name_or_path, **config_args, cache_dir=cache_dir)
-        self.auto_model = AutoModel.from_pretrained(model_name_or_path, config=self.config, cache_dir=cache_dir, **model_args)
 
         self._lora_adaptations = self.config.lora_adaptations
         if (
@@ -65,7 +65,10 @@ class Transformer(nn.Module):
         self._adaptation_map = {
             name: idx for idx, name in enumerate(self._lora_adaptations)
         }
-        self._default_task = None
+
+        self.default_task = model_args.pop('default_task', None)
+
+        self.auto_model = AutoModel.from_pretrained(model_name_or_path, config=self.config, cache_dir=cache_dir, **model_args)
 
         if max_seq_length is not None and "model_max_length" not in tokenizer_args:
             tokenizer_args["model_max_length"] = max_seq_length