enable autotrain params for lm training

huggingface · Apr 29, 2023 · b708f9e · b708f9e
1 parent 13431fe
commit b708f9e
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 30 deletions.
diff --git a/src/autotrain/app.py b/src/autotrain/app.py
@@ -340,22 +340,12 @@ def app():  # username, valid_orgs):
     st.sidebar.markdown("### Parameters")
     if model_choice != "AutoTrain":
         # on_change reset st.session_stat["jobs"]
-        if task != "lm_training":
-            param_choice = st.sidebar.selectbox(
-                "Parameter Choice",
-                ["AutoTrain", "Manual"],
-                key="param_choice",
-                on_change=on_change_reset_jobs(),
-            )
-        else:
-            param_choice = st.sidebar.selectbox(
-                "Parameter Choice",
-                ["Manual", "AutoTrain"],
-                key="param_choice",
-                on_change=on_change_reset_jobs(),
-                index=0,
-                disabled=True,
-            )
+        param_choice = st.sidebar.selectbox(
+            "Parameter Choice",
+            ["AutoTrain", "Manual"],
+            key="param_choice",
+            on_change=on_change_reset_jobs(),
+        )
     else:
         param_choice = st.sidebar.selectbox(
             "Parameter Choice", ["AutoTrain", "Manual"], key="param_choice", index=0, disabled=True

diff --git a/src/autotrain/params.py b/src/autotrain/params.py
@@ -246,20 +246,26 @@ def _tabular_binary_classification(self):
         }
 
     def _lm_training(self):
-        return {
-            "learning_rate": LMLearningRate,
-            "optimizer": Optimizer,
-            "scheduler": Scheduler,
-            "train_batch_size": LMTrainBatchSize,
-            "num_train_epochs": LMEpochs,
-            "percentage_warmup": PercentageWarmup,
-            "gradient_accumulation_steps": GradientAccumulationSteps,
-            "weight_decay": WeightDecay,
-            "lora_r": LoraR,
-            "lora_alpha": LoraAlpha,
-            "lora_dropout": LoraDropout,
-            "training_type": LMTrainingType,
-        }
+        if self.param_choice == "manual":
+            return {
+                "learning_rate": LMLearningRate,
+                "optimizer": Optimizer,
+                "scheduler": Scheduler,
+                "train_batch_size": LMTrainBatchSize,
+                "num_train_epochs": LMEpochs,
+                "percentage_warmup": PercentageWarmup,
+                "gradient_accumulation_steps": GradientAccumulationSteps,
+                "weight_decay": WeightDecay,
+                "lora_r": LoraR,
+                "lora_alpha": LoraAlpha,
+                "lora_dropout": LoraDropout,
+                "training_type": LMTrainingType,
+            }
+        if self.param_choice == "autotrain":
+            return {
+                "num_models": NumModels,
+            }
+        raise ValueError("param_choice must be either autotrain or manual")
 
     def _tabular_multi_class_classification(self):
         return self._tabular_binary_classification()