""" Configuration management and argument parsing for ML trainers. This module provides comprehensive configuration management including file-based configuration, command-line argument parsing, and validation with support for inheritance and environment variable override. """ import os import sys import json import yaml import argparse import logging from pathlib import Path from typing import Any, Dict, List, Optional, Union, Type from dataclasses import dataclass, field, fields, asdict from abc import ABC, abstractmethod import torch from .base import TrainerConfig, ModelFormat, TrainingState class ConfigError(Exception): """Exception raised for configuration-related errors.""" pass class ConfigValidator(ABC): """Abstract base class for configuration validators.""" @abstractmethod def validate(self, config: Dict[str, Any]) -> List[str]: """Validate configuration and return list of error messages.""" pass class TrainerConfigValidator(ConfigValidator): """Validator for trainer configuration.""" def validate(self, config: Dict[str, Any]) -> List[str]: """Validate trainer configuration.""" errors = [] # Required fields required_fields = ['trainer_name', 'model_name', 'data_dir'] for field in required_fields: if field not in config or not config[field]: errors.append(f"Required field '{field}' is missing or empty") # Numeric validations numeric_validations = { 'batch_size': (1, 1024), 'learning_rate': (1e-8, 1.0), 'num_epochs': (1, 10000), 'min_epochs': (0, 10000), 'max_epochs': (1, 50000), 'early_stopping_patience': (1, 1000), 'validation_split': (0.0, 1.0), 'test_split': (0.0, 1.0), 'sample_rate': (8000, 48000), 'n_mels': (10, 128), 'audio_length': (0.1, 30.0) } for field, (min_val, max_val) in numeric_validations.items(): if field in config: value = config[field] if not isinstance(value, (int, float)) or not (min_val <= value <= max_val): errors.append(f"Field '{field}' must be between {min_val} and {max_val}") # Validation split + test split check if 'validation_split' in config and 'test_split' in config: if config['validation_split'] + config['test_split'] >= 1.0: errors.append("validation_split + test_split must be less than 1.0") # Cross-field epoch validations if 'min_epochs' in config and 'num_epochs' in config: if config['min_epochs'] > config['num_epochs']: errors.append("min_epochs cannot be greater than num_epochs") if 'max_epochs' in config and 'num_epochs' in config: if config['max_epochs'] < config['num_epochs']: errors.append("max_epochs cannot be less than num_epochs") if 'min_epochs' in config and 'max_epochs' in config: if config['min_epochs'] > config['max_epochs']: errors.append("min_epochs cannot be greater than max_epochs") # Directory existence checks directory_fields = ['data_dir', 'output_dir', 'config_dir'] for field in directory_fields: if field in config and config[field]: path = Path(config[field]) if field == 'data_dir' and not path.exists(): errors.append(f"Data directory does not exist: {config[field]}") # Model format validation if 'model_format' in config: format_value = config['model_format'] if isinstance(format_value, str): valid_formats = [fmt.value for fmt in ModelFormat] if format_value not in valid_formats: errors.append(f"Invalid model format: {format_value}. Valid options: {valid_formats}") # Device validation if 'device' in config: device = config['device'] if device not in ['auto', 'cpu', 'cuda', 'mps']: errors.append(f"Invalid device: {device}. Valid options: auto, cpu, cuda, mps") return errors class WakewordConfigValidator(ConfigValidator): """Validator for wakeword trainer configuration.""" def validate(self, config: Dict[str, Any]) -> List[str]: """Validate wakeword-specific configuration.""" errors = [] # Wakeword-specific validations if 'wakeword_classes' in config: classes = config['wakeword_classes'] if not isinstance(classes, dict): errors.append("wakeword_classes must be a dictionary") elif len(classes) < 2: errors.append("At least 2 wakeword classes required (including negative)") # Audio processing for wakewords if 'audio_length' in config: if config['audio_length'] < 0.5 or config['audio_length'] > 5.0: errors.append("Wakeword audio_length should be between 0.5 and 5.0 seconds") return errors class VoiceRecognitionConfigValidator(ConfigValidator): """Validator for voice recognition trainer configuration.""" def validate(self, config: Dict[str, Any]) -> List[str]: """Validate voice recognition specific configuration.""" errors = [] # Speaker-specific validations if 'num_speakers' in config: num_speakers = config['num_speakers'] if not isinstance(num_speakers, int) or num_speakers < 2: errors.append("num_speakers must be at least 2") elif num_speakers > 1000: errors.append("num_speakers should not exceed 1000 for practical purposes") if 'embedding_dimension' in config: embedding_dim = config['embedding_dimension'] if not isinstance(embedding_dim, int) or embedding_dim < 64: errors.append("embedding_dimension must be at least 64") elif embedding_dim > 2048: errors.append("embedding_dimension should not exceed 2048") # Audio processing for voice recognition if 'audio_length' in config: if config['audio_length'] < 1.0 or config['audio_length'] > 10.0: errors.append("Voice recognition audio_length should be between 1.0 and 10.0 seconds") return errors class ConfigManager: """ Comprehensive configuration manager for ML trainers. Supports loading from files, command-line arguments, environment variables, and provides validation and inheritance capabilities. """ def __init__(self, config_dir: str = "./config", logger: Optional[logging.Logger] = None): """ Initialize configuration manager. Args: config_dir: Directory containing configuration files logger: Optional logger instance """ self.config_dir = Path(config_dir) self.logger = logger or logging.getLogger(__name__) self.validators: Dict[str, ConfigValidator] = { 'base': TrainerConfigValidator(), 'wakeword': WakewordConfigValidator(), 'voice_recognition': VoiceRecognitionConfigValidator() } self.loaded_configs: Dict[str, Dict[str, Any]] = {} def load_config_file(self, filepath: Union[str, Path]) -> Dict[str, Any]: """ Load configuration from file (JSON or YAML). Args: filepath: Path to configuration file Returns: Configuration dictionary """ filepath = Path(filepath) if not filepath.exists(): raise ConfigError(f"Configuration file not found: {filepath}") try: with open(filepath, 'r', encoding='utf-8') as f: if filepath.suffix.lower() in ['.yml', '.yaml']: try: config = yaml.safe_load(f) except ImportError: raise ConfigError("PyYAML not available for YAML configuration files") elif filepath.suffix.lower() == '.json': config = json.load(f) else: raise ConfigError(f"Unsupported configuration file format: {filepath.suffix}") self.logger.info(f"Loaded configuration from: {filepath}") return config or {} except (json.JSONDecodeError, yaml.YAMLError) as e: raise ConfigError(f"Error parsing configuration file {filepath}: {str(e)}") except Exception as e: raise ConfigError(f"Error loading configuration file {filepath}: {str(e)}") def save_config_file(self, config: Dict[str, Any], filepath: Union[str, Path]): """ Save configuration to file. Args: config: Configuration dictionary filepath: Path to save configuration """ filepath = Path(filepath) filepath.parent.mkdir(parents=True, exist_ok=True) try: with open(filepath, 'w', encoding='utf-8') as f: if filepath.suffix.lower() in ['.yml', '.yaml']: try: yaml.safe_dump(config, f, default_flow_style=False, allow_unicode=True) except ImportError: raise ConfigError("PyYAML not available for YAML configuration files") elif filepath.suffix.lower() == '.json': json.dump(config, f, indent=2, ensure_ascii=False) else: raise ConfigError(f"Unsupported configuration file format: {filepath.suffix}") self.logger.info(f"Saved configuration to: {filepath}") except Exception as e: raise ConfigError(f"Error saving configuration file {filepath}: {str(e)}") def load_trainer_config(self, trainer_name: str, config_name: str = None) -> TrainerConfig: """ Load trainer configuration with inheritance and validation. Args: trainer_name: Name of the trainer (wakeword, voice_recognition, etc.) config_name: Specific configuration name (defaults to trainer_name) Returns: TrainerConfig instance """ if config_name is None: config_name = trainer_name # Load base configuration base_config = self._load_base_config() # Load trainer-specific configuration trainer_config = self._load_trainer_specific_config(trainer_name, config_name) # Merge configurations (trainer-specific overrides base) merged_config = self._merge_configs(base_config, trainer_config) # Apply environment variable overrides merged_config = self._apply_env_overrides(merged_config) # Validate configuration self._validate_config(merged_config, trainer_name) # Convert to TrainerConfig object return self._dict_to_trainer_config(merged_config) def _load_base_config(self) -> Dict[str, Any]: """Load base trainer configuration.""" base_config_path = self.config_dir / "trainer_base.json" if base_config_path.exists(): return self.load_config_file(base_config_path) else: # Return default base configuration self.logger.info("Using default base configuration") return { "batch_size": 32, "learning_rate": 0.001, "num_epochs": 100, "min_epochs": 10, "early_stopping_patience": 10, "validation_split": 0.2, "test_split": 0.1, "sample_rate": 16000, "audio_length": 1.5, "n_mels": 40, "use_augmentation": True, "log_level": "INFO", "device": "auto" } def _load_trainer_specific_config(self, trainer_name: str, config_name: str) -> Dict[str, Any]: """Load trainer-specific configuration.""" config_paths = [ self.config_dir / f"{config_name}.json", self.config_dir / f"{config_name}.yml", self.config_dir / f"{config_name}.yaml", self.config_dir / f"{trainer_name}.json", self.config_dir / f"{trainer_name}.yml", self.config_dir / f"{trainer_name}.yaml" ] for config_path in config_paths: if config_path.exists(): return self.load_config_file(config_path) self.logger.warning(f"No specific configuration found for {trainer_name}/{config_name}") return {} def _merge_configs(self, base: Dict[str, Any], override: Dict[str, Any]) -> Dict[str, Any]: """Merge configuration dictionaries (override takes precedence).""" merged = base.copy() for key, value in override.items(): if key in merged and isinstance(merged[key], dict) and isinstance(value, dict): # Recursively merge nested dictionaries merged[key] = self._merge_configs(merged[key], value) else: merged[key] = value return merged def _apply_env_overrides(self, config: Dict[str, Any]) -> Dict[str, Any]: """Apply environment variable overrides.""" env_prefix = "TRIXY_TRAINER_" for key, value in os.environ.items(): if key.startswith(env_prefix): config_key = key[len(env_prefix):].lower() # Convert environment variable value to appropriate type try: # Try to parse as JSON first (for complex types) parsed_value = json.loads(value) except json.JSONDecodeError: # Fall back to string/numeric parsing if value.lower() in ['true', 'false']: parsed_value = value.lower() == 'true' elif value.isdigit(): parsed_value = int(value) else: try: parsed_value = float(value) except ValueError: parsed_value = value config[config_key] = parsed_value self.logger.info(f"Applied environment override: {config_key} = {parsed_value}") return config def _validate_config(self, config: Dict[str, Any], trainer_name: str): """Validate configuration using appropriate validators.""" all_errors = [] # Always validate with base validator base_errors = self.validators['base'].validate(config) all_errors.extend(base_errors) # Validate with trainer-specific validator if available if trainer_name in self.validators: trainer_errors = self.validators[trainer_name].validate(config) all_errors.extend(trainer_errors) if all_errors: error_msg = f"Configuration validation failed:\n" + "\n".join(f"- {error}" for error in all_errors) raise ConfigError(error_msg) def _dict_to_trainer_config(self, config_dict: Dict[str, Any]) -> TrainerConfig: """Convert configuration dictionary to TrainerConfig object.""" # Handle enum conversions if 'model_format' in config_dict: if isinstance(config_dict['model_format'], str): config_dict['model_format'] = ModelFormat(config_dict['model_format']) # Filter only valid TrainerConfig fields trainer_config_fields = {f.name for f in fields(TrainerConfig)} filtered_config = {k: v for k, v in config_dict.items() if k in trainer_config_fields} # Store additional config in custom_params custom_params = {k: v for k, v in config_dict.items() if k not in trainer_config_fields} if custom_params: # Merge with existing custom_params if any existing_custom_params = filtered_config.get('custom_params', {}) existing_custom_params.update(custom_params) filtered_config['custom_params'] = existing_custom_params try: return TrainerConfig(**filtered_config) except Exception as e: raise ConfigError(f"Error creating TrainerConfig: {str(e)}") def create_argument_parser(self, trainer_name: str = "trainer") -> argparse.ArgumentParser: """ Create command-line argument parser for trainer. Args: trainer_name: Name of the trainer for help text Returns: Configured ArgumentParser """ parser = argparse.ArgumentParser( description=f"Trixy {trainer_name.title()} Trainer", formatter_class=argparse.ArgumentDefaultsHelpFormatter ) # General arguments parser.add_argument( '--config', '-c', type=str, help='Configuration file path' ) parser.add_argument( '--trainer-name', type=str, default=trainer_name, help='Name of the trainer' ) parser.add_argument( '--model-name', type=str, help='Name of the model' ) parser.add_argument( '--output-dir', type=str, default='./models', help='Output directory for models' ) parser.add_argument( '--data-dir', type=str, default='./trainer/data', help='Data directory' ) # Training parameters training_group = parser.add_argument_group('Training Parameters') training_group.add_argument( '--batch-size', type=int, default=32, help='Batch size for training' ) training_group.add_argument( '--learning-rate', '--lr', type=float, default=0.001, help='Learning rate' ) training_group.add_argument( '--num-epochs', '--epochs', type=int, default=100, help='Number of training epochs' ) training_group.add_argument( '--min-epochs', type=int, default=10, help='Minimum number of epochs' ) training_group.add_argument( '--max-epochs', type=int, default=1000, help='Maximum number of epochs' ) training_group.add_argument( '--early-stopping-patience', type=int, default=10, help='Early stopping patience' ) # Model parameters model_group = parser.add_argument_group('Model Parameters') model_group.add_argument( '--model-format', choices=['pth', 'pt', 'onnx'], default='pth', help='Model file format' ) model_group.add_argument( '--use-mixed-precision', action='store_true', default=True, help='Use mixed precision training' ) model_group.add_argument( '--no-mixed-precision', dest='use_mixed_precision', action='store_false', help='Disable mixed precision training' ) # Data parameters data_group = parser.add_argument_group('Data Parameters') data_group.add_argument( '--validation-split', type=float, default=0.2, help='Validation split ratio' ) data_group.add_argument( '--test-split', type=float, default=0.1, help='Test split ratio' ) data_group.add_argument( '--sample-rate', type=int, default=16000, help='Audio sample rate' ) data_group.add_argument( '--audio-length', type=float, default=1.5, help='Audio length in seconds' ) data_group.add_argument( '--n-mels', type=int, default=40, help='Number of mel filterbanks' ) # Augmentation aug_group = parser.add_argument_group('Data Augmentation') aug_group.add_argument( '--use-augmentation', action='store_true', default=True, help='Use data augmentation' ) aug_group.add_argument( '--no-augmentation', dest='use_augmentation', action='store_false', help='Disable data augmentation' ) aug_group.add_argument( '--noise-factor', type=float, default=0.1, help='Noise augmentation factor' ) aug_group.add_argument( '--speed-factor', type=float, default=0.1, help='Speed augmentation factor' ) # Security security_group = parser.add_argument_group('Security') security_group.add_argument( '--password', type=str, help='Password for model encryption' ) security_group.add_argument( '--no-password-protection', dest='use_password_protection', action='store_false', default=True, help='Disable password protection' ) # Hardware hardware_group = parser.add_argument_group('Hardware') hardware_group.add_argument( '--device', choices=['auto', 'cpu', 'cuda', 'mps'], default='auto', help='Training device' ) hardware_group.add_argument( '--num-workers', type=int, default=4, help='Number of data loader workers' ) # Logging and checkpoints logging_group = parser.add_argument_group('Logging and Checkpoints') logging_group.add_argument( '--log-level', choices=['DEBUG', 'INFO', 'WARNING', 'ERROR'], default='INFO', help='Logging level' ) logging_group.add_argument( '--save-checkpoints', action='store_true', default=True, help='Save training checkpoints' ) logging_group.add_argument( '--no-checkpoints', dest='save_checkpoints', action='store_false', help='Disable checkpoint saving' ) logging_group.add_argument( '--checkpoint-interval', type=int, default=10, help='Checkpoint save interval (epochs)' ) # Resume training resume_group = parser.add_argument_group('Resume Training') resume_group.add_argument( '--resume', '--resume-from-checkpoint', type=str, help='Resume training from checkpoint' ) # Actions action_group = parser.add_argument_group('Actions') action_group.add_argument( '--train', action='store_true', default=True, help='Train the model' ) action_group.add_argument( '--test-only', action='store_true', help='Only test the model (requires trained model)' ) action_group.add_argument( '--prepare-data-only', action='store_true', help='Only prepare data without training' ) action_group.add_argument( '--validate-config-only', action='store_true', help='Only validate configuration' ) return parser def parse_args_to_config(self, args: argparse.Namespace, trainer_name: str) -> TrainerConfig: """ Convert command-line arguments to TrainerConfig. Args: args: Parsed command-line arguments trainer_name: Name of the trainer Returns: TrainerConfig instance """ # Start with configuration file if provided if hasattr(args, 'config') and args.config: config = self.load_config_file(args.config) else: config = self.load_trainer_config(trainer_name).to_dict() # Override with command-line arguments arg_dict = vars(args) for key, value in arg_dict.items(): if value is not None: # Convert hyphens to underscores for config keys config_key = key.replace('-', '_') config[config_key] = value # Handle special cases if hasattr(args, 'resume') and args.resume: config['resume_from_checkpoint'] = True config['checkpoint_path'] = args.resume if hasattr(args, 'model_format') and args.model_format: config['model_format'] = f".{args.model_format}" # Validate and convert to TrainerConfig self._validate_config(config, trainer_name) return self._dict_to_trainer_config(config) def generate_config_template(self, trainer_name: str, filepath: str = None) -> str: """ Generate configuration file template. Args: trainer_name: Name of the trainer filepath: Optional file path to save template Returns: Configuration template as string """ template = { "# Configuration for Trixy ML Trainer": None, "trainer_name": trainer_name, "model_name": f"{trainer_name}_model", "description": f"Configuration for {trainer_name} trainer", "# Data settings": None, "data_dir": f"./trainer/data/{trainer_name}", "output_dir": "./models", "validation_split": 0.2, "test_split": 0.1, "# Training parameters": None, "batch_size": 32, "learning_rate": 0.001, "num_epochs": 100, "min_epochs": 10, "early_stopping_patience": 10, "# Model settings": None, "model_format": ".pth", "use_mixed_precision": True, "use_password_protection": True, "# Audio processing": None, "sample_rate": 16000, "audio_length": 1.5, "n_mels": 40, "n_fft": 512, "hop_length": 160, "# Data augmentation": None, "use_augmentation": True, "noise_factor": 0.1, "speed_factor": 0.1, "pitch_factor": 0.1, "volume_factor": 0.2, "# Hardware": None, "device": "auto", "num_workers": 4, "# Logging": None, "log_level": "INFO", "save_checkpoints": True, "checkpoint_interval": 10 } # Add trainer-specific settings if trainer_name == "wakeword": template.update({ "# Wakeword specific": None, "wakeword_classes": { 0: "custom", 1: "system_command", 2: "negative" }, "detection_threshold": 0.5 }) elif trainer_name == "voice_recognition": template.update({ "# Voice recognition specific": None, "num_speakers": 2, "embedding_dimension": 256, "architecture": "ecapa_tdnn" }) # Convert to JSON string template_str = "{\n" for key, value in template.items(): if value is None: # Comment line template_str += f" \"{key}\",\n" elif isinstance(value, str): template_str += f" \"{key}\": \"{value}\",\n" elif isinstance(value, dict): template_str += f" \"{key}\": {json.dumps(value, indent=4)},\n" else: template_str += f" \"{key}\": {json.dumps(value)},\n" template_str = template_str.rstrip(',\n') + "\n}" # Save to file if requested if filepath: with open(filepath, 'w', encoding='utf-8') as f: f.write(template_str) self.logger.info(f"Configuration template saved to: {filepath}") return template_str def list_available_configs(self) -> List[str]: """List all available configuration files.""" configs = [] if self.config_dir.exists(): for pattern in ['*.json', '*.yml', '*.yaml']: configs.extend([f.stem for f in self.config_dir.glob(pattern)]) return sorted(list(set(configs))) def get_config_info(self, config_name: str) -> Dict[str, Any]: """Get information about a configuration file.""" config_paths = [ self.config_dir / f"{config_name}.json", self.config_dir / f"{config_name}.yml", self.config_dir / f"{config_name}.yaml" ] for config_path in config_paths: if config_path.exists(): config = self.load_config_file(config_path) return { 'name': config_name, 'path': str(config_path), 'format': config_path.suffix, 'size': config_path.stat().st_size, 'trainer_name': config.get('trainer_name', 'unknown'), 'model_name': config.get('model_name', 'unknown'), 'description': config.get('description', ''), 'keys': list(config.keys()) } raise ConfigError(f"Configuration not found: {config_name}") def create_trainer_cli(trainer_name: str) -> argparse.ArgumentParser: """ Create command-line interface for a trainer. Args: trainer_name: Name of the trainer Returns: Configured ArgumentParser """ config_manager = ConfigManager() return config_manager.create_argument_parser(trainer_name) def load_config_from_args(args: argparse.Namespace, trainer_name: str) -> TrainerConfig: """ Load configuration from command-line arguments. Args: args: Parsed command-line arguments trainer_name: Name of the trainer Returns: TrainerConfig instance """ config_manager = ConfigManager() return config_manager.parse_args_to_config(args, trainer_name)