TrixyZ/trainer/components/callbacks/training_callbacks.py

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
    ))