Fix bugs in Optuna integration with the prompt2model demo script (#374)

* Fix bugs in Optuna integration

* Lint

Author: viswavi

prompt2model/utils/config.py

@@ -2,8 +2,8 @@
 MAX_SUPPORTED_BATCH_SIZE = 4
 
 DEFAULT_HYPERPARAMETERS_SPACE = {
-    "min_num_train_epochs": 10,
-    "max_num_train_epochs": 20,
+    "min_num_train_epochs": 5,
+    "max_num_train_epochs": 15,
     "save_strategy": ["no"],
     "evaluation_strategy": ["no"],
     "per_device_train_batch_size": MAX_SUPPORTED_BATCH_SIZE,

prompt2model_demo.py

@@ -91,12 +91,12 @@ def parse_model_size_limit(line: str, default_size=3e9) -> float:
         return default_size
     model_units = {"B": 1e0, "KB": 1e3, "MB": 1e6, "GB": 1e9, "TB": 1e12, "PB": 1e15}
     unit_disambiguations = {
-        "B": ["b", "bytes"],
         "KB": ["Kb", "kb", "kilobytes"],
         "MB": ["Mb", "mb", "megabytes"],
         "GB": ["Gb", "gb", "gigabytes"],
         "TB": ["Tb", "tb", "terabytes"],
         "PB": ["Pb", "pb", "petabytes"],
+        "B": ["b", "bytes"],
     }
     unit_matched = False
     for unit, disambiguations in unit_disambiguations.items():
@@ -110,6 +110,7 @@ def parse_model_size_limit(line: str, default_size=3e9) -> float:
         numerical_part = line.strip()[: -len(unit_name)].strip()
     else:
         numerical_part = line.strip()
+
     if not str.isdecimal(numerical_part):
         raise ValueError(
             "Invalid input. Please enter a number (integer " + "or number with units)."
@@ -368,28 +369,40 @@ def main():
 
         if line == "y":
             line_print("Starting training with hyperparameter selection.")
+            default_min_num_epochs = DEFAULT_HYPERPARAMETERS_SPACE[
+                "min_num_train_epochs"
+            ]
+            min_num_epochs = input(
+                f"Enter min number of epochs. Press enter to use default value ({default_min_num_epochs}): "  # noqa E501
+            )
+            default_max_num_epochs = DEFAULT_HYPERPARAMETERS_SPACE[
+                "max_num_train_epochs"
+            ]
             max_num_epochs = input(
-                "Enter max number of epochs: Press enter to use default value: "
+                f"Enter max number of epochs. Press enter to use default value ({default_max_num_epochs}): "  # noqa E501
             )
+            default_num_trials = 10
             num_trials = input(
-                "Enter the number of trials to conduct hypeparamter search. Press enter to use default value: "  # noqa E501
+                f"Enter the number of trials (maximum number of hyperparameter configurations to consider) for hyperparameter search. Press enter to use default value ({default_num_trials}): "  # noqa E501
             )
             default_batch_size = DEFAULT_HYPERPARAMETERS_SPACE[
                 "per_device_train_batch_size"
             ]  # noqa E501
             max_batch_size = input(
                 "Enter the max batch size. "
-                + f"Press enter to use default: {default_batch_size}"
+                + f"Press enter to use default ({default_batch_size}): "
             )
 
+            min_num_epochs = (
+                default_min_num_epochs if min_num_epochs == "" else eval(min_num_epochs)
+            )
             max_num_epochs = (
-                DEFAULT_HYPERPARAMETERS_SPACE["max_num_train_epochs"]
-                if max_num_epochs == ""
-                else eval(max_num_epochs)
+                default_max_num_epochs if max_num_epochs == "" else eval(max_num_epochs)
             )
             num_trials = 1 if num_trials == "" else eval(num_trials)
+
             max_batch_size = (
-                DEFAULT_HYPERPARAMETERS_SPACE["per_device_train_batch_size"][0]
+                DEFAULT_HYPERPARAMETERS_SPACE["per_device_train_batch_size"]
                 if max_batch_size == ""
                 else eval(max_batch_size)
             )
@@ -406,12 +419,13 @@ def main():
             line_print("Starting training.")
 
             trained_model, trained_tokenizer = OptunaParamSelector(
-                n_trial=num_trials,
+                n_trials=num_trials,
                 trainer=trainer,
             ).select_from_hyperparameters(
                 training_datasets=training_datasets,
                 validation=validation_datasets,
                 hyperparameters={
+                    "min_num_train_epochs": min_num_epochs,
                     "max_num_train_epochs": max_num_epochs,
                     "per_device_train_batch_size": [max_batch_size],
                 },