import logging from dash.dependencies import Input as DashInput, Output as DashOutput, State as DashState from dash import callback_context, dcc from utils.file_utils import ensure_directory_exists, list_files_in_directory, update_training_time from threading import Thread import time import os import importlib.util import json from utils.training_interface import TrainerInterface # Configure logging to file logger = logging.getLogger(__name__) logging.basicConfig(filename='trainer.log', encoding='utf-8', level=logging.DEBUG) TRAINING_HOTWORD_DIR = './training/hotword/' MODEL_DIR = './model/hotword/' BACKGROUND_DIR = './training/background/' SAMPLER_DIR = './training/sampler/' training_status = "Not started" training_start_time = None training_thread = None training_info = {} # Global variables for progress tracking current_progress = 0 current_loss = 0.0 current_accuracy = 0.0 total_epochs = 0 total_batches = 0 def load_trainer_modules(): trainer_modules = {} trainer_dir = './trainer_scripts/' for file_name in os.listdir(trainer_dir): if file_name.endswith('.py'): module_name = file_name[:-3] file_path = os.path.join(trainer_dir, file_name) spec = importlib.util.spec_from_file_location(module_name, file_path) module = importlib.util.module_from_spec(spec) spec.loader.exec_module(module) for attr_name in dir(module): attr = getattr(module, attr_name) if isinstance(attr, type) and issubclass(attr, TrainerInterface) and attr != TrainerInterface: trainer_modules[module_name] = attr return trainer_modules trainer_modules = load_trainer_modules() def register_training_callbacks(app): @app.callback( [DashOutput("start-training-hotword", "disabled"), DashOutput("optimize-training-hotword", "disabled"), DashOutput("training-status", "children"), DashOutput("training-timer", "children"), DashOutput("training-chart", "figure"), DashOutput("training-progress", "value")], [DashInput("start-training-hotword", "n_clicks"), DashInput("optimize-training-hotword", "n_clicks"), DashInput("stop-training-hotword", "n_clicks"), DashInput('interval-component', 'n_intervals')], [DashState('url', 'pathname'), DashState("training-method-dropdown", "value"), DashState("dropout-rate", "value"), DashState("num-epochs", "value"), DashState("batch-size", "value")] ) def manage_training(n_clicks_start, n_clicks_optimize, n_clicks_stop, n_intervals, pathname, training_method, dropout_rate, num_epochs, batch_size): global training_status, training_start_time, training_thread, training_info, logger global current_progress, current_loss, current_accuracy, total_epochs, total_batches model_name = pathname.split('/')[-1] ctx = callback_context if not os.path.exists(MODEL_DIR + model_name): os.mkdir(MODEL_DIR + model_name) if not ctx.triggered: logger.error("Callback context triggered is empty.") return False, False, f"Status: {training_status}", "", {}, 0 triggered_id = ctx.triggered[0]['prop_id'].split('.')[0] if triggered_id == "start-training-hotword" and n_clicks_start: training_status = "Preprocessing" training_start_time = time.time() training_info[model_name] = {"accuracy": [], "loss": [], "epochs": 0} if training_method in trainer_modules: trainer_class = trainer_modules[training_method] trainer = trainer_class() else: logger.error(f"Unsupported training method: {training_method}") return False, False, f"Status: {training_status}", "", {}, 0 trainer.set_model_path(MODEL_DIR + model_name + "/") trainer.set_training_directories([TRAINING_HOTWORD_DIR + model_name + "/", TRAINING_HOTWORD_DIR + "systemcommand/"]) trainer.set_sample_directory(SAMPLER_DIR) trainer.set_background_directory(BACKGROUND_DIR) trainer.set_dropout_rate(dropout_rate) trainer.set_epochs(num_epochs) trainer.set_batch_size(batch_size) training_thread = Thread(target=trainer.start_new) training_thread.start() logger.warning(f"Started training thread for model: {model_name}") return True, True, f"Status: {training_status}", f"Started at: {time.ctime(training_start_time)}", {}, 0 elif triggered_id == "optimize-training-hotword" and n_clicks_optimize: training_status = "Optimizing" training_start_time = time.time() training_info[model_name] = {"accuracy": [], "loss": [], "epochs": 0} if training_method in trainer_modules: trainer_class = trainer_modules[training_method] trainer = trainer_class() else: logger.error(f"Unsupported training method: {training_method}") return False, False, f"Status: {training_status}", "", {}, 0 trainer.set_model_path(MODEL_DIR + model_name + "/") trainer.set_training_directories([TRAINING_HOTWORD_DIR + model_name + "/", TRAINING_HOTWORD_DIR + "systemcommand/"]) trainer.set_sample_directory(SAMPLER_DIR) trainer.set_background_directory(BACKGROUND_DIR) trainer.set_dropout_rate(dropout_rate) trainer.set_epochs(num_epochs) trainer.set_batch_size(batch_size) training_thread = Thread(target=trainer.start) training_thread.start() logger.warning(f"Started optimization thread for model: {model_name}") return True, True, f"Status: {training_status}", f"Started at: {time.ctime(training_start_time)}", {}, 0 elif triggered_id == "stop-training-hotword" and n_clicks_stop: training_status = "Stopped" logger.debug(f"Stopped training for model: {model_name}") return False, False, f"Status: {training_status}", "", {}, 0 elif triggered_id == "interval-component" and training_start_time and training_status in ["Training", "Optimizing"]: elapsed_time = time.time() - training_start_time figure = { "data": [ {"x": list(range(len(training_info[model_name]["accuracy"]))), "y": training_info[model_name]["accuracy"], "type": "line", "name": "Accuracy", "line": {"color": "blue"}}, {"x": list(range(len(training_info[model_name]["loss"]))), "y": training_info[model_name]["loss"], "type": "line", "name": "Loss", "line": {"color": "green"}}, ], "layout": { "title": "Training Progress", "xaxis": {"title": "Epoch"}, "yaxis": {"title": "Value", "range": [0, 2]}, "height": 200, "width": 500, "plot_bgcolor": "darkgray", "gridcolor": "lightgray" } } progress = current_progress training_info[model_name]["accuracy"].append(current_accuracy) training_info[model_name]["loss"].append(current_loss) training_info[model_name]["epochs"] = total_epochs return True, True, f"Status: {training_status}", f"Training time: {int(elapsed_time // 60)} minutes {int(elapsed_time % 60)} seconds", figure, progress training_data = list_files_in_directory(f"{TRAINING_HOTWORD_DIR}/{model_name}", extension='.wav') if training_status in ["Training", "Optimizing"] or not training_data: return True, True, f"Status: {training_status}", "", {}, 0 return False, False, f"Status: {training_status}", "", {}, 0 app.layout.children.append(dcc.Interval( id='interval-component', interval=1*1000, # in milliseconds n_intervals=0 ))