tests

Author: connermanuel

src/together/resources/finetune.py

@@ -195,10 +195,10 @@ def create_finetune_request(
             lr_scheduler_args=LinearLRSchedulerArgs(min_lr_ratio=min_lr_ratio),
         )
 
-    training_method_cls: TrainingMethodSFT | TrainingMethodDPO = TrainingMethodSFT(
-        train_on_inputs=train_on_inputs
-    )
-    if training_method == "dpo":
+    training_method_cls: TrainingMethodSFT | TrainingMethodDPO
+    if training_method == "sft":
+        training_method_cls = TrainingMethodSFT(train_on_inputs=train_on_inputs)
+    elif training_method == "dpo":
         training_method_cls = TrainingMethodDPO(dpo_beta=dpo_beta)
 
     finetune_request = FinetuneRequest(

tests/unit/test_finetune_resources.py

@@ -281,3 +281,32 @@ def test_bad_training_method():
             training_file=_TRAINING_FILE,
             training_method="NON_SFT",
         )
+
+
+@pytest.mark.parametrize("train_on_inputs", [True, False, "auto", None])
+def test_train_on_inputs_for_sft(train_on_inputs):
+    request = create_finetune_request(
+        model_limits=_MODEL_LIMITS,
+        model=_MODEL_NAME,
+        training_file=_TRAINING_FILE,
+        training_method="sft",
+        train_on_inputs=train_on_inputs,
+    )
+    assert request.training_method.method == "sft"
+    if isinstance(train_on_inputs, bool):
+        assert request.training_method.train_on_inputs is train_on_inputs
+    else:
+        assert request.training_method.train_on_inputs == "auto"
+
+
+def test_train_on_inputs_not_supported_for_dpo():
+    with pytest.raises(
+        ValueError, match="train_on_inputs is only supported for SFT training"
+    ):
+        _ = create_finetune_request(
+            model_limits=_MODEL_LIMITS,
+            model=_MODEL_NAME,
+            training_file=_TRAINING_FILE,
+            training_method="dpo",
+            train_on_inputs=True,
+        )