anyreach-ai
/

dualturn-endpointing

+"""
+Configuration management for turn-taking model training.
+Provides dataclass-based config with validation and YAML loading.
+"""
+from dataclasses import dataclass, field, asdict
+from pathlib import Path
+from typing import Literal, Optional, List
+import yaml
+@dataclass
+class ModelConfig:
+    """Model architecture hyperparameters."""
+    # Qwen backbone
+    qwen_model_name: str = "Qwen/Qwen2.5-0.5B"
+    # LoRA config
+    lora_r: 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"]
+    )
+    # Audio embeddings
+    num_codebooks: int = 8
+    codebook_size: int = 2049
+    hidden_dim: int = 896
+    # Depth predictor
+    depth_dim: int = 512
+    # Backbone type: "qwen" (default LLM) or "lstm" (lightweight baseline)
+    backbone_type: str = "qwen"
+    # LSTM backbone config (only used when backbone_type="lstm")
+    lstm_hidden_dim: int = 512
+    lstm_num_layers: int = 2
+    lstm_dropout: float = 0.1
+    lstm_bidirectional: bool = True
+    # Transformer backbone config (only used when backbone_type="transformer")
+    transformer_num_layers: int = 10
+    transformer_num_heads: int = 14
+    transformer_ff_dim: int = 3584
+    transformer_dropout: float = 0.1
+    # Layer probing: None = last layer (default), "weighted" = learned weighted avg,
+    # int = specific layer index (0=embedding, 1-24=transformer layers)
+    probe_layer: Optional[str] = None
+    # Random init: if True, initialize Qwen with random weights instead of pretrained
+    random_init: bool = False
+    # Input mode: "discrete" (codebook embeddings) or "continuous" (Mimi encoder features)
+    input_mode: str = "discrete"
+    mimi_feat_dim: int = 512
+    # AudioAdapter (dual-task only)
+    use_audio_adapter: bool = False
+    audio_adapter_heads: int = 8
+    audio_adapter_ff_dim: int = 2048
+    audio_adapter_dropout: float = 0.1
+    # Prediction head architecture: "linear" (legacy) or "mlp" (2-layer MLP for sparse tasks)
+    head_type: str = "linear"
+    head_hidden_dim: int = 256
+    head_dropout: float = 0.1
+    # Per-task learned layer attention (ELMo-style, per task)
+    per_task_layer_attention: bool = False
+    # Per-codebook output heads in depth predictor (vs shared MLP)
+    per_codebook_heads: bool = False
+    # Two-stream backbone (autoresearch best_model): siamese per-channel + cross-channel block.
+    # Requires backbone_type="transformer" and input_mode="continuous".
+    two_stream_backbone: bool = False
+@dataclass
+class DataConfig:
+    """Data loading and augmentation hyperparameters."""
+    # Paths
+    processed_dir: str = "data/otospeech_processed_npy"
+    splits_path: str = "data/splits.json"
+    metadata_path: str = "data/otospeech_metadata.json"
+    # Windowing
+    window_frames: int = 125  # 10s at 12.5Hz
+    hop_frames_train: int = 25  # 2s hop for training (overlapping)
+    hop_frames_val: int = 125  # 10s hop for val/test (non-overlapping)
+    # Augmentation
+    channel_swap_prob: float = 0.5  # Channel swap augmentation (train only)
+    # Soft labels
+    soft_labels: bool = True
+    sigma_before: float = 3.0  # frames (240ms at 12.5Hz)
+    sigma_after: float = 1.0   # frames (80ms at 12.5Hz)
+    # Balanced sampling
+    balanced_sampling: bool = True
+    shift_oversample_ratio: float = 3.0  # Oversample windows with shifts
+    dynamic_window_sampling: bool = False  # Random window positions around shifts (Stage 3)
+    # Event label widening: dilate sparse labels to N frames (1 = no widening)
+    event_label_width: int = 1
+    # Dual-task data paths (SWB + Oto separate)
+    swb_processed_dir: Optional[str] = None
+    swb_splits_path: Optional[str] = None
+    oto_processed_dir: Optional[str] = None
+    oto_splits_path: Optional[str] = None
+    oto_metadata_path: Optional[str] = None
+    max_text_tokens: int = 512
+    # S2S dataset (production call recordings) — included as extra dataset alongside Oto/SWB
+    s2s_processed_dir: Optional[str] = None
+    s2s_splits_path: Optional[str] = None
+@dataclass
+class TrainingConfig:
+    """Training loop hyperparameters."""
+    # Stage
+    stage: Literal["stage2", "stage2_dual_task", "stage3", "linear_probe", "full_finetune", "lstm_baseline"] = "stage2"
+    # Optimization
+    learning_rate: float = 1e-4
+    weight_decay: float = 0.01
+    adam_beta1: float = 0.9
+    adam_beta2: float = 0.999
+    adam_eps: float = 1e-8
+    max_grad_norm: float = 1.0  # Gradient clipping
+    # Scheduler
+    warmup_steps: int = 1000
+    scheduler_type: str = "cosine"  # "cosine", "linear", or "constant"
+    # Training loop
+    batch_size: int = 16
+    num_epochs: int = 5
+    max_steps: Optional[int] = None  # If set, overrides num_epochs
+    gradient_accumulation_steps: int = 1
+    # Loss weights -- context-aware tasks (v3)
+    weight_eot: float = 1.0        # End of Turn
+    weight_hold: float = 1.0       # Turn Hold / Pause
+    weight_bot: float = 1.0        # Beginning of Turn
+    weight_bc: float = 1.0         # Backchannel
+    weight_vad: float = 1.0        # Voice Activity Detection (BCE)
+    weight_codebook: float = 0.0   # Codebook prediction loss weight
+    # Sparse event loss type: "focal" (recommended) or "wbce" (weighted BCE)
+    # Focal loss avoids gradient spikes from high pos_weight by down-weighting
+    # easy negatives adaptively. Use "wbce" only for backward compat.
+    sparse_loss_type: str = "focal"
+    # Focal loss per-task alpha (positive class weight, in [0,1]).
+    # Higher alpha = more weight on positives. Combined with gamma, this is
+    # much gentler than wBCE pos_weight while achieving better gradient balance.
+    eot_alpha: float = 0.75
+    hold_alpha: float = 0.60
+    bot_alpha: float = 0.80
+    bc_alpha: float = 0.80
+    focal_gamma_sparse: float = 2.0
+    # Per-task pos_weight for wBCE (only used when sparse_loss_type="wbce").
+    # Derived from combined dataset positive rates after 5-frame widening + VAD mask:
+    #   EOT~2.7%, HOLD~5.6%, BOT~1.3%, BC~1.4%
+    eot_pos_weight: float = 37.0
+    hold_pos_weight: float = 17.0
+    bot_pos_weight: float = 75.0
+    bc_pos_weight: float = 73.0
+    sparse_pos_weight: float = 20.0  # legacy fallback if per-task not set in config
+    # Legacy weights (backward compat -- set to 0.0 for new training)
+    weight_shift: float = 0.0
+    weight_end: float = 0.0
+    weight_start: float = 0.0
+    focal_gamma: float = 2.0
+    focal_alpha: float = 0.75
+    # Future VAD projection -- VAP-style binned voice activity prediction
+    weight_fvad: float = 0.0    # Future VAD projection loss weight (0 = disabled)
+    fvad_bins: List[int] = field(default_factory=lambda: [3, 6, 12, 25])
+    # Bin edges in frames at 12.5Hz: [3,6,12,25] ->
+    #   bin0: t+1..t+3  (80-240ms),  bin1: t+4..t+6  (320-480ms)
+    #   bin2: t+7..t+12 (560-960ms),  bin3: t+13..t+25 (1.04-2.0s)
+    # Text loss (dual-task)
+    weight_text: float = 0.0         # ASR text prediction loss weight (Stage 2 dual-task)
+    mode: str = "codebook"           # forward mode: "codebook", "dual_task", "shift"
+    # Full finetune: merge LoRA into base and unfreeze all params
+    full_finetune: bool = False
+    # Per-codebook weighting: [code0, ..., code7]. Default equal.
+    # Moshi-inspired: code0 carries prosody -> weight higher.
+    codebook_weights: Optional[List[float]] = None
+    # Future VAD auxiliary task (legacy, kept for old config compat)
+    vad_lookahead_frames: int = 3
+    # Label mode
+    label_mode: str = "start_end"  # "shift" (legacy) or "start_end" (v2)
+    # Validation & checkpointing
+    eval_every_steps: int = 500
+    save_every_steps: int = 1000
+    early_stopping_patience: int = 5  # Stop after N evals without improvement
+    max_val_batches: Optional[int] = None  # Limit val batches (None = full val set; useful for quick tests)
+    save_generated_audio: bool = False     # Save autoregressive audio samples each validation (Stage-1)
+    gen_context_frames: int = 125          # Context frames for generation (125 = 10s at 12.5Hz)
+    gen_gen_frames: int = 62              # Frames to generate (62 = ~5s at 12.5Hz)
+    # System
+    num_workers: int = 4
+    pin_memory: bool = True
+    mixed_precision: bool = False  # bf16 training (A100 native support)
+    seed: int = 42
+    # Logging
+    log_every_steps: int = 10
+    log_perplexity: bool = True  # Log perplexity for codebook prediction (Stage 2)
+    wandb_project: Optional[str] = None  # "turn-taking-interspeech"
+    wandb_run_name: Optional[str] = None
+    experiment_name: str = "default"
+    # Paths
+    checkpoint_dir: str = "checkpoints"
+    log_dir: str = "logs"
+    stage2_checkpoint: Optional[str] = None  # Path to Stage 2 checkpoint (for Stage 3)
+@dataclass
+class ExperimentConfig:
+    """Complete experiment configuration."""
+    model: ModelConfig = field(default_factory=ModelConfig)
+    data: DataConfig = field(default_factory=DataConfig)
+    training: TrainingConfig = field(default_factory=TrainingConfig)
+    @classmethod
+    def from_yaml(cls, path: str) -> "ExperimentConfig":
+        """Load config from YAML file."""
+        with open(path, 'r') as f:
+            data = yaml.safe_load(f)
+        config = cls(
+            model=ModelConfig(**data.get("model", {})),
+            data=DataConfig(**data.get("data", {})),
+            training=TrainingConfig(**data.get("training", {})),
+        )
+        return config
+    def to_yaml(self, path: str):
+        """Save config to YAML file."""
+        data = {
+            "model": asdict(self.model),
+            "data": asdict(self.data),
+            "training": asdict(self.training),
+        }
+        with open(path, 'w') as f:
+            yaml.dump(data, f, default_flow_style=False, sort_keys=False)
+    def to_dict(self) -> dict:
+        """Convert to dictionary."""
+        return {
+            "model": asdict(self.model),
+            "data": asdict(self.data),
+            "training": asdict(self.training),
+        }