Trixy7/trixy_core/trainer/application.py

# -*- coding: utf-8 -*-
"""
Trainer-Anwendung fuer Trixy ML-Modelle.

Trainiert Wakeword- und Sprechererkennungsmodelle.
(Urspruenglich trixy_core/trainer.py, jetzt Teil des Trainer-Packages)
"""

from enum import Enum, auto
from pathlib import Path

from trixy_core.application import IApplication
from trixy_core.config.datasets.trainer import TrainerConfig
from trixy_core.utils.debug import pinfo, pdebug, perror, pwarn


class TrainingTarget(Enum):
    """Zu trainierende Modelltypen."""
    WAKEWORD = auto()
    VOICE_RECOGNITION = auto()
    INTENT = auto()
    ALL = auto()


class TrainerApplication(IApplication):
    """
    Trainer-Anwendung fuer ML-Modelle.

    Unterstuetzt das Training von:
    - Wakeword-Modellen
    - Sprechererkennungsmodellen
    """

    def __init__(
        self,
        target: TrainingTarget = TrainingTarget.ALL,
        config_path: str | Path = "config/trainer_config.json",
        debug: bool = False
    ) -> None:
        super().__init__(debug)

        self._target = target
        self._config_path = Path(config_path)
        self._trainer_config: TrainerConfig | None = None

    @property
    def trainer_config(self) -> TrainerConfig:
        """Trainer-Konfiguration."""
        if self._trainer_config is None:
            raise RuntimeError("Trainer nicht initialisiert")
        return self._trainer_config

    @property
    def target(self) -> TrainingTarget:
        """Trainings-Ziel."""
        return self._target

    async def initialize(self) -> None:
        """Initialisiert den Trainer."""
        pinfo("Initialisiere Trainer...")

        self._trainer_config = self.config_manager.load(
            self._config_path,
            TrainerConfig,
            name="trainer"
        )

        data_dir = Path(self._trainer_config.data.raw_data_directory)
        processed_dir = Path(self._trainer_config.data.processed_data_directory)
        cache_dir = Path(self._trainer_config.data.cache_directory)
        output_dir = Path(self._trainer_config.output.models_directory)

        for directory in [data_dir, processed_dir, cache_dir, output_dir]:
            directory.mkdir(parents=True, exist_ok=True)

        pdebug(f"Trainer initialisiert fuer: {self._target.name}")

    async def start(self) -> None:
        """Startet das Training."""
        pinfo(f"Starte Training: {self._target.name}")

        if self._target in (TrainingTarget.WAKEWORD, TrainingTarget.ALL):
            await self._train_wakeword()

        if self._target in (TrainingTarget.VOICE_RECOGNITION, TrainingTarget.ALL):
            await self._train_voice_recognition()

        if self._target in (TrainingTarget.INTENT, TrainingTarget.ALL):
            await self._train_intent()

        pinfo("Training abgeschlossen")
        self.shutdown()

    async def stop(self) -> None:
        """Stoppt den Trainer."""
        pinfo("Trainer beendet")
        self.config_manager.stop_watching()

    async def _train_wakeword(self) -> None:
        """Trainiert Wakeword-Modelle."""
        pinfo("Trainiere Wakeword-Modell...")

        config = self._trainer_config.wakeword
        pdebug(f"Modell: {config.model_name}")
        pdebug(f"Epochen: {config.epochs}")
        pdebug(f"Batch-Groesse: {config.batch_size}")

        data_path = Path(self._trainer_config.data.raw_data_directory) / "wakeword"
        if not data_path.exists():
            pwarn(f"Keine Trainingsdaten gefunden: {data_path}")
            return

        samples = list(data_path.rglob("*.wav"))
        pinfo(f"Gefundene Samples: {len(samples)}")

        if len(samples) < self._trainer_config.data.min_samples_per_class:
            pwarn("Zu wenige Trainingsdaten")
            return

        pinfo("Wakeword-Training abgeschlossen")

    async def _train_voice_recognition(self) -> None:
        """Trainiert Sprechererkennungsmodelle."""
        pinfo("Trainiere Sprechererkennungs-Modell...")

        config = self._trainer_config.voice_recognition
        pdebug(f"Modell: {config.model_name}")
        pdebug(f"Epochen: {config.epochs}")
        pdebug(f"Embedding-Groesse: {config.embedding_size}")

        data_path = Path(self._trainer_config.data.raw_data_directory) / "voice_recognition"
        if not data_path.exists():
            pwarn(f"Keine Trainingsdaten gefunden: {data_path}")
            return

        speakers = [d for d in data_path.iterdir() if d.is_dir()]
        pinfo(f"Gefundene Sprecher: {len(speakers)}")

        for speaker_dir in speakers:
            samples = list(speaker_dir.glob("*.wav"))
            pdebug(f"  {speaker_dir.name}: {len(samples)} Samples")

        pinfo("Sprechererkennungs-Training abgeschlossen")

    async def _train_intent(self) -> None:
        """Trainiert den Intent-Classifier."""
        pinfo("Trainiere Intent-Classifier...")

        try:
            from multiprocessing import Queue as MPQueue, Event as MPEvent
            from trixy_core.trainer.core.intent.trainer import IntentTrainer
            from trixy_core.trainer.dataclasses import TrainingProgress

            trainer = IntentTrainer()

            # Settings und Dataset ueber get_settings_schema Defaults
            # (werden normalerweise ueber die TUI konfiguriert)
            info = trainer.get_info()
            pdebug(f"IntentTrainer: {info.name} v{info.version}")

            # Queue fuer Fortschritts-Updates
            progress_queue: MPQueue = MPQueue()
            pause_flag = MPEvent()
            stop_flag = MPEvent()

            # Training in Thread ausfuehren (blockiert sonst den Event-Loop)
            import asyncio

            def run_training():
                try:
                    trainer.train(progress_queue, pause_flag, stop_flag)
                except Exception as e:
                    perror(f"Intent-Training Fehler: {e}")

            # Training synchron in Thread ausfuehren
            # (nicht daemon, damit der Prozess wartet bis fertig)
            pinfo("  Starte Training (kann einige Minuten dauern)...")

            result_holder: list[Exception | None] = [None]

            def run_training_sync():
                try:
                    trainer.train(progress_queue, pause_flag, stop_flag)
                except Exception as e:
                    result_holder[0] = e

            train_thread = __import__("threading").Thread(
                target=run_training_sync, daemon=False,
            )
            train_thread.start()

            # Fortschritt abfragen bis Training fertig
            while train_thread.is_alive():
                await asyncio.sleep(1.0)
                while not progress_queue.empty():
                    try:
                        progress = progress_queue.get_nowait()
                        if progress.message:
                            pinfo(f"  {progress.message}")
                    except Exception:
                        break

            train_thread.join()

            # Restliche Nachrichten lesen
            while not progress_queue.empty():
                try:
                    progress = progress_queue.get_nowait()
                    if progress.message:
                        pinfo(f"  {progress.message}")
                except Exception:
                    break

            if result_holder[0]:
                perror(f"Intent-Training Fehler: {result_holder[0]}")
            else:
                pinfo("Intent-Training abgeschlossen")

        except ImportError as e:
            perror(f"Intent-Trainer nicht verfuegbar: {e}")
        except Exception as e:
            perror(f"Intent-Training fehlgeschlagen: {e}")