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."""
        import yaml
        from pathlib import Path as _Path

        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."""

    datasets_dir: str = "/home/vscode/datasets"
    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)
    synthetic_mix: SyntheticMixConfig = Field(default_factory=SyntheticMixConfig)
    input_filter: FilterConfig = Field(default_factory=FilterConfig)
    target_filter: FilterConfig = Field(default_factory=FilterConfig)

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

    embedding_dim: int = 16
    latent_width: int = 128
    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"

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

    learning_rate: float = 1e-3
    lr_schedule: LrScheduleConfig = Field(default_factory=LrScheduleConfig)
    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 = 20
    tflite_rep_batches: int = 8
    input_bit_depth: int = 16
    band_metrics: BandMetricsConfig = Field(default_factory=BandMetricsConfig)
    physiokit_metrics: PhysiokitMetricsConfig = Field(default_factory=PhysiokitMetricsConfig)

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)