helena 0.1.0

//! The raw experiment-config schema: plain serde data, shape-proof only.
//!
//! Every type here is a `Raw*` document tier. The top-level document and the
//! sections that carry no free-form map use `#[serde(deny_unknown_fields)]`, so
//! a typo'd or misplaced key is a parse error — the property that lets the CLI
//! drop its mirrored strict-config schema. The three *component* sections
//! ([`RawDataEncoder`], [`RawAudioCodec`], [`RawGenerator`]) instead carry a
//! deliberate flattened `extra: Metadata` map, so `deny_unknown_fields` cannot
//! apply to them (serde flatten captures the surplus keys); their unknown-key
//! discipline is enforced downstream by
//! [`reject_unknown_extra`](super::reject_unknown_extra), which the resolvers
//! that *consume* `extra` call with their recognized set.

use serde::{Deserialize, Serialize};

use crate::error::Result;
use crate::latent::{KindTag, Metadata};
use crate::manifest::SourceKind;

/// A raw experiment configuration as deserialized, before resolution.
///
/// The only path from this document to the validated
/// [`ExperimentSpec`](super::ExperimentSpec) is
/// [`resolve`](Self::resolve); this type carries no invariants of its own.
///
/// **Wire-format note.** `source_data.type` is now a [`SourceKind`] enum and
/// `audio_codec.latent_type` a [`KindTag`] — both were free strings on the
/// prior schema. The old `audio_codec.latent_type = "rvq"` (and the
/// `"rvq_tokens"` variant) is *gone*: tokens-with-codebooks is the single token
/// model now, so a config naming `rvq` no longer parses.
#[derive(Clone, Debug, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct RawExperimentConfig {
    /// Project namespace the experiment belongs to.
    pub project: String,
    /// Experiment name, unique within the project.
    pub experiment_name: String,
    /// Target audio sample rate, in hertz.
    pub sample_rate: u32,
    /// Nominal clip length the pipeline trains and generates at.
    pub clip_duration_seconds: f32,
    /// Where source items live and how they align to target audio.
    pub source_data: RawSourceData,
    /// Where the training audio lives and how it is windowed.
    pub target_audio: RawTargetAudio,
    /// Which encoder turns source data into conditioning latents.
    pub data_encoder: RawDataEncoder,
    /// Which codec maps between waveform and audio latents.
    pub audio_codec: RawAudioCodec,
    /// Which generator maps conditioning latents to audio latents.
    pub conditional_generator: RawGenerator,
    /// Optimizer and schedule settings for the training run.
    pub training: RawTraining,
    /// Default generation knobs recorded with the config.
    #[serde(default)]
    pub generation: RawGeneration,
}

/// How source data relates to target audio during training.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum Alignment {
    /// Every source item has its own aligned target clip.
    Paired,
    /// Source and audio correspond as sets, not item by item.
    WeaklyPaired,
    /// No correspondence; source and audio are independent corpora.
    Unpaired,
}

/// Where source items live and how they relate to target audio.
#[derive(Clone, Debug, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct RawSourceData {
    /// Source modality (a [`SourceKind`], parsed once — no longer a stringly
    /// `kind`).
    #[serde(rename = "type")]
    pub kind: SourceKind,
    /// Filesystem path to the source dataset.
    pub path: String,
    /// How source items correspond to target audio.
    #[serde(default = "default_alignment")]
    pub alignment: Alignment,
}

fn default_alignment() -> Alignment {
    Alignment::Paired
}

/// Where the training audio lives and how it is segmented.
#[derive(Clone, Debug, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct RawTargetAudio {
    /// Filesystem path to the target-audio dataset.
    pub path: String,
    /// How clips are windowed for training.
    #[serde(default)]
    pub segmentation: RawSegmentation,
}

/// Audio windowing for training.
#[derive(Clone, Copy, Debug, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct RawSegmentation {
    /// Length of each training window, in seconds.
    pub window_seconds: f32,
    /// Overlap between consecutive windows, in seconds.
    #[serde(default)]
    pub overlap_seconds: f32,
}

impl Default for RawSegmentation {
    fn default() -> Self {
        Self {
            window_seconds: 10.0,
            overlap_seconds: 0.0,
        }
    }
}

/// Selects and configures the data encoder.
///
/// Carries a flattened `extra` map, so — unlike the leaf sections — it cannot
/// use `deny_unknown_fields`; the resolver validates `extra` against its
/// recognized set via [`reject_unknown_extra`](super::reject_unknown_extra).
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct RawDataEncoder {
    /// Encoder plugin key.
    #[serde(rename = "type")]
    pub kind: String,
    /// Width of each conditioning latent the encoder emits.
    pub latent_dim: usize,
    /// Encoder-specific extra settings (consumed by the resolver).
    #[serde(flatten, default)]
    pub extra: Metadata,
}

/// Selects and configures the audio codec.
///
/// Carries a flattened `extra` map (see [`RawDataEncoder`]).
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct RawAudioCodec {
    /// Codec plugin key.
    #[serde(rename = "type")]
    pub kind: String,
    /// Latent kind the codec produces: a [`KindTag`] (`"continuous"` /
    /// `"tokens"`). Formerly a free `latent_type` string whose `"rvq"` value is
    /// gone.
    pub latent_type: KindTag,
    /// Whether the codec is frozen during training.
    #[serde(default = "default_true")]
    pub freeze: bool,
    /// Codec-specific extra settings (consumed by the resolver).
    #[serde(flatten, default)]
    pub extra: Metadata,
}

fn default_true() -> bool {
    true
}

/// Selects and configures the conditional generator.
///
/// Carries a flattened `extra` map (see [`RawDataEncoder`]).
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct RawGenerator {
    /// Generator plugin key.
    #[serde(rename = "type")]
    pub kind: String,
    /// Architecture-specific settings (consumed by the resolver).
    #[serde(flatten, default)]
    pub extra: Metadata,
}

/// Optimizer and schedule settings for a training run.
#[derive(Clone, Debug, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct RawTraining {
    /// Number of items per optimizer step.
    pub batch_size: usize,
    /// Optimizer learning rate.
    pub learning_rate: f64,
    /// Total optimizer steps to run.
    pub max_steps: u64,
    /// Seed for reproducible training.
    #[serde(default)]
    pub seed: u64,
}

/// Default generation knobs recorded in the config.
///
/// Left inference-only: these knobs shape sampling, not trained weights, so
/// they are excluded from the [`Recipe`](super::Recipe) that gates checkpoint
/// reload.
#[derive(Clone, Debug, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct RawGeneration {
    /// Sampling temperature; `0.0` is deterministic.
    pub temperature: f32,
    /// Classifier-free guidance strength; `1.0` applies conditioning as-is.
    pub guidance_scale: f32,
    /// How many candidates to generate per request.
    pub num_candidates: u32,
}

impl Default for RawGeneration {
    fn default() -> Self {
        Self {
            temperature: 1.0,
            guidance_scale: 1.0,
            num_candidates: 1,
        }
    }
}

impl RawExperimentConfig {
    /// Parse a config from a JSON string.
    ///
    /// Proves the *document shape* only — the field set, their types, and the
    /// `deny_unknown_fields` gate. Semantic invariants (positive rates and
    /// durations, non-empty component keys, advancing segmentation) are
    /// [`resolve`](Self::resolve)'s responsibility. A shape failure (malformed
    /// JSON or a missing/mistyped/surplus field) is an [`Error::Json`]; a
    /// semantic failure is an [`Error::Validation`], so the two stay
    /// distinguishable to a caller.
    ///
    /// [`Error::Json`]: crate::Error::Json
    /// [`Error::Validation`]: crate::Error::Validation
    pub fn from_json_str(s: &str) -> Result<Self> {
        Ok(serde_json::from_str(s)?)
    }

    /// Parse a config from a TOML string.
    ///
    /// Like [`from_json_str`](Self::from_json_str), this proves document shape
    /// only; call [`resolve`](Self::resolve) for semantic invariants. A shape
    /// failure is an [`Error::Toml`], kept distinct from the
    /// [`Error::Validation`] a semantic check returns.
    ///
    /// [`Error::Toml`]: crate::Error::Toml
    /// [`Error::Validation`]: crate::Error::Validation
    pub fn from_toml_str(s: &str) -> Result<Self> {
        Ok(toml::from_str(s)?)
    }

    /// Serialize the config to pretty JSON.
    pub fn to_json_string(&self) -> Result<String> {
        Ok(serde_json::to_string_pretty(self)?)
    }
}