"""Configuration for the prosody-to-text pipeline."""

from __future__ import annotations

from transformers import PretrainedConfig


class ProsodyToTextConfig(PretrainedConfig):
    """Everything needed to rebuild the pipeline.

    The mel contract is the part that matters. `mel_bins` and `length_ratio`
    must match what the conformer was trained with: the loader falls back to
    defaults when they are absent, and a wrong `length_ratio` mis-times the
    mel against the audio without failing, which is the worst way for it to
    be wrong.
    """

    model_type = "prosody_to_text"

    def __init__(
        self,
        # --- the conformer's mel contract -------------------------------
        mel_bins: int = 12,
        length_ratio: float = 0.5,
        conformer_subfolder: str = "conformer",
        # --- feature extraction -----------------------------------------
        # 200 Hz, and the physical-range normalisation the conformer was
        # trained with. These mirror the training preprocessor; changing one
        # without the others silently moves the model off its own
        # distribution.
        feature_rate_hz: float = 200.0,
        f0_scale: float = 500.0,
        intensity_scale: float = 100.0,
        # --- the decoder -------------------------------------------------
        whisper_subfolder: str = "whisper",
        tokenizer_subfolder: str = "tokenizer",
        # Ten, matching the published inference runs. Wider was tried and
        # made the transcriptions worse: on a signal this weak a wider beam
        # searches harder for the highest-likelihood sequence, which is the
        # caption boilerplate the checkpoint overlearned.
        num_beams: int = 10,
        language: str = "en",
        no_repeat_ngram_size: int = 3,
        max_new_tokens: int = 225,
        # On, matching the dataset runs the published numbers come from: ban
        # numerals, and take the token budget from the utterance's own
        # duration rather than a flat limit. Set False to decode
        # unconstrained, which is what the drawing demo does.
        constrain: bool = True,
        **kwargs,
    ) -> None:
        self.mel_bins = mel_bins
        self.length_ratio = length_ratio
        self.conformer_subfolder = conformer_subfolder

        self.feature_rate_hz = feature_rate_hz
        self.f0_scale = f0_scale
        self.intensity_scale = intensity_scale

        self.whisper_subfolder = whisper_subfolder
        self.tokenizer_subfolder = tokenizer_subfolder
        self.num_beams = num_beams
        self.language = language
        self.no_repeat_ngram_size = no_repeat_ngram_size
        self.max_new_tokens = max_new_tokens
        self.constrain = constrain

        super().__init__(**kwargs)