"""Training configuration dataclass for lr_whisper. All hyperparameters are read from environment variables via ``TrainingConfig.from_env()``. Use ``get_config()`` to obtain a process-wide singleton (lazy, never called at import time). """ import os from dataclasses import dataclass, field from typing import List, Optional from utils import env_flag, normalize_train_mode @dataclass class TrainingConfig: # Model / dataset model_id: str = "Qwen/Qwen3-ASR-0.6B" dataset_name: str = "openslr/librispeech_asr" dataset_config: str = "clean" # Training batch / epoch batch_size: int = 4 grad_accum_steps: int = 4 num_epochs: int = 10 eval_samples: int = 1000 pretrain_eval_samples: int = 10 # Latent tokens n_latent: int = 4 is_deq_train: bool = False deq_train_iters: int = 15 halt_threshold: float = 0.0 latent_drop_prob: float = 0.15 latent_input_noise_std: float = 0.5 value_forced_neg_prob: float = 0.3 latent_use_bounded_delta: bool = True latent_use_injection_gate: bool = True latent_use_embedding_anchor: bool = True train_max_samples: int = 0 checkpoint_prefix: str = "" # Thought layout thought_mode: str = "prefix" # "prefix" | "interleaved" thought_group_size: int = 1 # words per NT token in interleaved mode # Training mode train_mode: str = "latent" # Learning rates lr_adapter: float = 1e-4 lr_scale: float = 5e-5 lr_baseline_ft: float = 1e-5 lr_prompt_tuning: float = 5e-4 lr_lora_r16: float = 1e-4 # Prompt text user_prompt_text: str = "Transcribe the audio into text." # Prompt-tuning hyperparameters prompt_tuning_init_mode: str = "text" prompt_tuning_num_virtual_tokens: int = 4 prompt_tuning_init_text: str = "Transcribe the audio into text." # LoRA hyperparameters lora_rank: int = 16 lora_alpha: int = 32 lora_dropout: float = 0.05 lora_target_modules: List[str] = field( default_factory=lambda: [ "q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj", ] ) # Delta stabilisation / scale delta_tanh_c: float = 5.0 scale_max: float = 3.0 scale_init: float = 0.2 # Regularisation weights w_smooth: float = 0.0 # Disabled — replaced by w_cycle w_cycle: float = 0.1 # Cycle consistency loss (vocabulary-independent) - main regularizer w_value: float = 3.0 # Logging intervals log_every: int = 50 grad_log_every: int = 50 # Derived flags (set in __post_init__) use_latent_reasoning: bool = False use_prompt_tuning: bool = False use_lora_r16: bool = False use_peft_mode: bool = False def __post_init__(self) -> None: valid_modes = {"latent", "baseline", "prompt_tuning", "lora_r16"} if self.train_mode not in valid_modes: raise ValueError( f"Unsupported train_mode={self.train_mode!r}. " f"Use one of: {sorted(valid_modes)}" ) valid_thought_modes = {"prefix", "interleaved"} if self.thought_mode not in valid_thought_modes: raise ValueError( f"Unsupported thought_mode={self.thought_mode!r}. " f"Use one of: {sorted(valid_thought_modes)}" ) if self.thought_group_size < 1: raise ValueError(f"thought_group_size must be >= 1, got {self.thought_group_size}") self.use_latent_reasoning = self.train_mode == "latent" self.use_prompt_tuning = self.train_mode == "prompt_tuning" self.use_lora_r16 = self.train_mode == "lora_r16" self.use_peft_mode = self.use_lora_r16 @classmethod def from_env(cls) -> "TrainingConfig": """Construct a ``TrainingConfig`` by reading environment variables.""" lora_modules_raw = os.getenv( "LORA_TARGET_MODULES", "q_proj,k_proj,v_proj,o_proj,gate_proj,up_proj,down_proj", ) lora_target_modules = [x.strip() for x in lora_modules_raw.split(",") if x.strip()] train_mode = normalize_train_mode(os.getenv("TRAIN_MODE", "latent")) default_dataset = "openslr/librispeech_asr" default_config = "clean" if train_mode == "latent": default_dataset = "SpeechTest/extreme_asr_pony" default_config = "default" return cls( model_id=os.getenv("MODEL_ID", "Qwen/Qwen3-ASR-0.6B"), dataset_name=os.getenv("DATASET_NAME", default_dataset), dataset_config=os.getenv("DATASET_CONFIG", default_config), batch_size=int(os.getenv("BATCH_SIZE", "4")), grad_accum_steps=int(os.getenv("GRAD_ACCUM_STEPS", "4")), num_epochs=int(os.getenv("NUM_EPOCHS", "10")), eval_samples=int(os.getenv("EVAL_SAMPLES", "1000")), pretrain_eval_samples=int(os.getenv("PRETRAIN_EVAL_SAMPLES", "10")), n_latent=int(os.getenv("N_LATENT", "4")), is_deq_train=os.getenv("IS_DEQ_TRAIN", "true").strip().lower() == "true", deq_train_iters=int(os.getenv("DEQ_TRAIN_ITERS", "15")), halt_threshold=float(os.getenv("HALT_THRESHOLD", "0.0")), latent_drop_prob=float(os.getenv("LATENT_DROP_PROB", "0.15")), latent_input_noise_std=float(os.getenv("LATENT_INPUT_NOISE_STD", "0.5")), value_forced_neg_prob=float(os.getenv("VALUE_FORCED_NEG_PROB", "0.3")), latent_use_bounded_delta=env_flag("LATENT_USE_BOUNDED_DELTA", True), latent_use_injection_gate=env_flag("LATENT_USE_INJECTION_GATE", True), latent_use_embedding_anchor=env_flag("LATENT_USE_EMBEDDING_ANCHOR", True), train_max_samples=int(os.getenv("TRAIN_MAX_SAMPLES", "0")), checkpoint_prefix=os.getenv("CHECKPOINT_PREFIX", "").strip(), thought_mode=os.getenv("THOUGHT_MODE", "prefix").strip().lower(), thought_group_size=int(os.getenv("THOUGHT_GROUP_SIZE", "1")), train_mode=train_mode, lr_adapter=float(os.getenv("LR_ADAPTER", "1e-4")), lr_scale=float(os.getenv("LR_SCALE", "5e-5")), lr_baseline_ft=float(os.getenv("LR_BASELINE_FT", "1e-5")), lr_prompt_tuning=float(os.getenv("LR_PROMPT_TUNING", "5e-4")), lr_lora_r16=float(os.getenv("LR_LORA_R16", "1e-4")), user_prompt_text=os.getenv( "USER_PROMPT_TEXT", "Transcribe the audio into text." ).strip(), prompt_tuning_init_mode=os.getenv("PROMPT_TUNING_INIT_MODE", "text").strip().lower(), prompt_tuning_num_virtual_tokens=int( os.getenv("PROMPT_TUNING_NUM_VIRTUAL_TOKENS", "4") ), prompt_tuning_init_text=os.getenv( "PROMPT_TUNING_INIT_TEXT", "Transcribe the audio into text." ).strip(), lora_rank=int(os.getenv("LORA_RANK", "16")), lora_alpha=int(os.getenv("LORA_ALPHA", "32")), lora_dropout=float(os.getenv("LORA_DROPOUT", "0.05")), lora_target_modules=lora_target_modules, delta_tanh_c=float(os.getenv("DELTA_TANH_C", "5.0")), scale_max=float(os.getenv("SCALE_MAX", "3.0")), scale_init=float(os.getenv("SCALE_INIT", "0.2")), w_smooth=float(os.getenv("W_SMOOTH", "0.0")), w_cycle=float(os.getenv("W_CYCLE", "0.1")), w_value=float(os.getenv("W_VALUE", "3.0")), log_every=int(os.getenv("LOG_EVERY", "50")), grad_log_every=int(os.getenv("GRAD_LOG_EVERY", "50")), ) _config: Optional[TrainingConfig] = None def get_config() -> TrainingConfig: """Return the process-wide ``TrainingConfig`` singleton (lazy init from env).""" global _config if _config is None: _config = TrainingConfig.from_env() return _config