allow for teacher to be passed in as instantiated model (#2170) (#2172)

Author: Sara Adkins

src/sparseml/transformers/finetune/model_args.py

@@ -30,6 +30,12 @@ class ModelArguments:
             )
         },
     )
+    distill_teacher: Optional[str] = field(
+        default=None,
+        metadata={
+            "help": "Teacher model (a trained text generation model)",
+        },
+    )
     config_name: Optional[str] = field(
         default=None,
         metadata={

src/sparseml/transformers/finetune/text_generation.py

@@ -155,10 +155,10 @@ def intialize_model_from_path(
     )
     teacher_config = (
         AutoConfig.from_pretrained(
-            training_args.distill_teacher,
+            model_args.distill_teacher,
             use_auth_token=True if model_args.use_auth_token else None,
         )
-        if training_args.distill_teacher
+        if model_args.distill_teacher
         else None
     )
 
@@ -208,11 +208,11 @@ def intialize_model_from_path(
 
     teacher = (
         SparseAutoModel.text_generation_from_pretrained(
-            model_name_or_path=training_args.distill_teacher,
+            model_name_or_path=model_args.distill_teacher,
             sequence_length=None,  # use model default
             **teacher_kwargs,
         )
-        if training_args.distill_teacher is not None
+        if model_args.distill_teacher is not None
         else None
     )
 
@@ -289,7 +289,7 @@ def main(
 
     # Detecting last checkpoint.
     last_checkpoint = None
-    teacher = None
+    teacher = model_args.distill_teacher
     model_path = None
     model = model_args.model
     # Load tokenizer

src/sparseml/transformers/finetune/training_args.py

@@ -32,12 +32,6 @@ class TrainingArguments(HFTrainingArgs):
         arguments
     """
 
-    distill_teacher: Optional[str] = field(
-        default=None,
-        metadata={
-            "help": "Teacher model (a trained text generation model)",
-        },
-    )
     best_model_after_epoch: int = field(
         default=None,
         metadata={"help": "Epoch after which best model will be saved."},