more fixes 20240228 (#1342) [skip ci]

* add missing evals_per_epoch setting

* more pydantic fixes

* more fixes

* move test from normalization to validation

* increase eval size for sample packing tests

src/axolotl/cli/__init__.py

@@ -13,7 +13,6 @@ from threading import Thread
 from typing import Any, Dict, List, Optional, Union
 from urllib.parse import urlparse
 
-import gradio as gr
 import requests
 import torch
 import yaml
@@ -215,6 +214,8 @@ def do_inference_gradio(
     cfg: DictDefault,
     cli_args: TrainerCliArgs,
 ):
+    import gradio as gr
+
     model, tokenizer = load_model_and_tokenizer(cfg=cfg, cli_args=cli_args)
     prompter = cli_args.prompter
     default_tokens = {"unk_token": "<unk>", "bos_token": "<s>", "eos_token": "</s>"}

src/axolotl/utils/config/__init__.py

@@ -164,9 +164,6 @@ def normalize_config(cfg):
         ]
     ) or cfg.is_qwen_derived_model
 
-    if isinstance(cfg.learning_rate, str):
-        cfg.learning_rate = float(cfg.learning_rate)
-
     if isinstance(cfg.pretraining_dataset, dict):
         cfg.pretraining_dataset = [cfg.pretraining_dataset]
 

src/axolotl/utils/config/models/input/v0_4_1/__init__.py

@@ -302,6 +302,13 @@ class HyperparametersConfig(BaseModel):
             )
         return batch_size
 
+    @field_validator("learning_rate")
+    @classmethod
+    def convert_learning_rate(cls, learning_rate):
+        if learning_rate and isinstance(learning_rate, str):
+            learning_rate = float(learning_rate)
+        return learning_rate
+
 
 class ModelOutputConfig(BaseModel):
     """model save configuration subset"""
@@ -368,6 +375,7 @@ class AxolotlInputConfig(
     rl: Optional[RLType] = None
 
     datasets: Optional[conlist(Union[SFTDataset, DPODataset], min_length=1)] = None  # type: ignore
+    test_datasets: Optional[conlist(Union[SFTDataset, DPODataset], min_length=1)] = None  # type: ignore
     dataset_prepared_path: Optional[str] = None
     dataset_shard_num: Optional[int] = None
     dataset_shard_idx: Optional[int] = None
@@ -456,6 +464,7 @@ class AxolotlInputConfig(
     warmup_steps: Optional[int] = None
     warmup_ratio: Optional[float] = None
     eval_steps: Optional[Union[int, float]] = None
+    evals_per_epoch: Optional[Union[int]] = None
     evaluation_strategy: Optional[str] = None
     save_steps: Optional[Union[int, float]] = None
     saves_per_epoch: Optional[int] = None
@@ -463,6 +472,7 @@ class AxolotlInputConfig(
     save_total_limit: Optional[int] = None
     logging_steps: Optional[int] = None
     early_stopping_patience: Optional[int] = None
+    load_best_model_at_end: Optional[bool] = False
 
     neftune_noise_alpha: Optional[float] = None
 

src/axolotl/utils/trainer.py

@@ -255,7 +255,7 @@ def calculate_total_num_steps(cfg, train_dataset, update=True):
                 train_dataset.remove_columns(["length"]),
                 batch_sampler=sampler,
             )
-            data_loader_len = len(data_loader) // batch_size
+            data_loader_len = len(data_loader) // cfg.batch_size
             actual_eff = sampler.efficiency()
             LOG.debug(f"data_loader_len: {data_loader_len}", main_process_only=True)
             # FIXME: is there a bug here somewhere? the total num steps depends