Configs¶

class PpgRvqConfig(BaseModel):
    """Top-level configuration for PPG RVQ training pipeline."""

    run_name: str = "ppg_rvq_run"
    data: DataConfig = Field(default_factory=DataConfig)
    model: ModelConfig = Field(default_factory=ModelConfig)
    training: TrainingConfig = Field(default_factory=TrainingConfig)
    evaluation: EvaluationConfig = Field(default_factory=EvaluationConfig)
    output: OutputConfig = Field(default_factory=OutputConfig)

    @classmethod
    def from_yaml(cls, path: str) -> PpgRvqConfig:
        """Load config from a YAML file, merged with defaults."""
        from pathlib import Path as _Path

        import yaml

        with _Path(path).open("r") as f:
            user_cfg = yaml.safe_load(f) or {}
        return cls.model_validate(user_cfg)

class DataConfig(BaseModel):
    """Data loading and preprocessing configuration."""

    dataset_id: str | None = Field(
        default=None,
        description="Identifier registered in compressionkit.datasets.contract (see #26).",
    )
    datasets_dir: str = Field(default_factory=default_datasets_dir)
    dataset_glob: str = "mesa-commercial-use/polysomnography/edfs/*.edf"
    sampling_rate: int = 64
    frame_size: int = 320
    segment_samples: int = 7680  # 64 Hz * 120 s
    offset_samples: int = 192  # 64 Hz * 3 s
    target_label: str = "Pleth"
    batch_size: int = 64
    buffer_size: int = 10_000
    steps_per_epoch: int = 250
    epochs: int = 100
    epsilon: float = 1e-3
    gaussian_noise: list[float] = Field(default_factory=lambda: [0.01, 0.1])
    shuffle_seed: int = 42
    streaming: StreamingConfig = Field(default_factory=StreamingConfig)
    cache: CacheConfig = Field(default_factory=CacheConfig)
    unified_cache: UnifiedCacheConfig = Field(default_factory=UnifiedCacheConfig)
    synthetic_mix: SyntheticMixConfig = Field(default_factory=SyntheticMixConfig)
    input_filter: FilterConfig = Field(default_factory=FilterConfig)
    target_filter: FilterConfig = Field(default_factory=FilterConfig)
    transform: TransformConfig = Field(default_factory=TransformConfig)
    augmentation: AugmentationConfig = Field(default_factory=AugmentationConfig)
    artifact_suite: ArtifactSuiteConfig = Field(default_factory=ArtifactSuiteConfig)

class ModelConfig(BaseModel):
    """RVQ autoencoder model architecture configuration."""

    embedding_dim: int = 16
    latent_width: int = 128
    codebook_sizes: list[int] | None = Field(
        default=None,
        description=(
            "Per-level codebook sizes for hierarchical / multi-rate RVQ. "
            "Length must equal num_levels. When set, overrides latent_width "
            "for the RVQ bottleneck (e.g. [512, 256, 128, 64] gives a "
            "coarse-to-fine hierarchy). When None, all levels use latent_width."
        ),
    )
    num_levels: int = 2
    num_stages: int = 4
    base_filters: int = 32
    multiplier: float = 1.25
    beta: float = 0.25
    encoder_block_norm: str = "batch"
    encoder_head_norm: str = "none"
    decoder_block_norm: str = "none"
    decoder_head_norm: str = "layer"
    use_ema: bool = False
    ema_decay: float = 0.99
    revive_dead_codes: bool = Field(
        default=False,
        description="Enable EnCodec/DAC-style dead-code revival in EMA RVQ.",
    )
    revive_threshold: float = Field(
        default=0.03,
        description=(
            "Codes whose normalized usage falls below revive_threshold/K are resampled from the current batch."
        ),
    )
    kmeans_init: bool = Field(
        default=False,
        description=(
            "Mini-batch k-means warm-start of the EMA RVQ codebooks before training. Run once eagerly via the trainer."
        ),
    )
    encoder_type: str = Field(
        default="default",
        description="Encoder architecture: 'default', 'inverted_residual', 'soundstream', 'mlp', or 'transformer'.",
    )
    decoder_type: str = Field(
        default="default",
        description="Decoder architecture: 'default', 'ssm', 'hierarchical', 'soundstream', 'mlp', or 'transformer'.",
    )
    encoder_blocks_per_stage: int = Field(
        default=1,
        description="Number of blocks/layers per stage (MLP hidden layers, transformer layers, etc.).",
    )
    bottleneck_type: str = Field(
        default="rvq",
        description="Bottleneck type: 'rvq', 'fsq', or 'significance'.",
    )

class TrainingConfig(BaseModel):
    """Training hyperparameters and schedule configuration."""

    learning_rate: float = 1e-3
    lr_schedule: LrScheduleConfig = Field(default_factory=LrScheduleConfig)
    derivative_loss: DerivativeLossConfig = Field(default_factory=DerivativeLossConfig)
    spectral_loss: SpectralLossConfig = Field(default_factory=SpectralLossConfig)
    val_metric: str = "mse"
    val_mode: str = "min"
    early_stop_patience: int = 25
    reduce_lr_on_plateau: bool = True
    reduce_lr_factor: float = 0.5
    reduce_lr_patience: int = 8
    reduce_lr_min_lr: float = 1e-5
    selection_metric: str = "val_mse"
    validation_steps: int | None = None

class EvaluationConfig(BaseModel):
    """Post-training evaluation configuration."""

    num_samples: int = 1000
    """How many random validation samples to evaluate (CSV + metrics)."""
    num_plot_samples: int = 50
    """Subset of ``num_samples`` that also receive a PNG plot artifact."""
    tflite_rep_batches: int = 8
    input_bit_depth: int = 16
    band_metrics: BandMetricsConfig = Field(default_factory=BandMetricsConfig)
    physiokit_metrics: PhysiokitMetricsConfig = Field(default_factory=PhysiokitMetricsConfig)
    long_recording: LongRecordingEvalConfig = Field(default_factory=LongRecordingEvalConfig)

class OutputConfig(BaseModel):
    """Output directory and logging configuration."""

    results_root: str = "results"
    log_file: str | None = None
    tensorboard: bool = True
    wandb: WandbConfig = Field(default_factory=WandbConfig)