add: support selecting saved configuration files and loading training parameters

This commit is contained in:
hzhaoy 2024-06-04 10:33:43 +08:00
parent c4f50865ad
commit b27c4cfcb3
2 changed files with 15 additions and 2 deletions

View File

@ -6,7 +6,7 @@ from ...extras.constants import TRAINING_STAGES
from ...extras.misc import get_device_count
from ...extras.packages import is_gradio_available
from ..common import DEFAULT_DATA_DIR, list_checkpoints, list_datasets
from ..utils import change_stage, check_output_dir, list_output_dirs
from ..utils import change_stage, check_output_dir, list_output_dirs, list_config_paths
from .data import create_preview_box
@ -259,7 +259,7 @@ def create_train_tab(engine: "Engine") -> Dict[str, "Component"]:
with gr.Row():
initial_dir = gr.Textbox(visible=False, interactive=False)
output_dir = gr.Dropdown(allow_custom_value=True)
config_path = gr.Textbox()
config_path = gr.Dropdown(allow_custom_value=True)
with gr.Row():
device_count = gr.Textbox(value=str(get_device_count() or 1), interactive=False)
@ -317,5 +317,6 @@ def create_train_tab(engine: "Engine") -> Dict[str, "Component"]:
output_dir.change(
list_output_dirs, [model_name, finetuning_type, initial_dir], [output_dir], concurrency_limit=None
).then(check_output_dir, inputs=[lang, model_name, finetuning_type, output_dir], concurrency_limit=None)
config_path.change(list_config_paths, outputs=[config_path], concurrency_limit=None)
return elem_dict

View File

@ -176,6 +176,18 @@ def list_output_dirs(model_name: str, finetuning_type: str, initial_dir: str) ->
return gr.Dropdown(choices=output_dirs)
def list_config_paths() -> "gr.Dropdown":
"""
Lists all the saved configuration files that can be loaded.
"""
if os.path.exists(DEFAULT_CONFIG_DIR) and os.path.isdir(DEFAULT_CONFIG_DIR):
config_files = [file_name for file_name in os.listdir(DEFAULT_CONFIG_DIR) if file_name.endswith(".yaml")]
else:
config_files = []
return gr.Dropdown(choices=config_files)
def check_output_dir(lang: str, model_name: str, finetuning_type: str, output_dir: str) -> None:
r"""
Check if output dir exists.