configuration_zuna.py

"""
ZunaConfig: HuggingFace PretrainedConfig wrapper for the Zyphra ZUNA foundation model.

Maps all fields from DecoderTransformerArgs (and its parent chain
BaseTransformerArgs -> DecoderArgs -> DecoderTransformerArgs) to a standard
HF config, so the model can be loaded with AutoConfig / trust_remote_code=True.

Field sources:
  BaseTransformerArgs  – lingua/transformer.py
  DecoderArgs          – xattn.py
  DecoderTransformerArgs – transformer.py (AY2latent_bci)
"""

from typing import List, Optional, Union
from transformers import PretrainedConfig


class ZunaConfig(PretrainedConfig):
    model_type = "zuna"

    def __init__(
        self,
        # ── BaseTransformerArgs ──────────────────────────────────────────────
        dim: int = 1024,
        n_layers: int = 10,
        head_dim: Optional[int] = None,
        n_heads: int = 8,
        n_kv_heads: Optional[int] = None,
        ffn_dim_multiplier: Optional[float] = None,
        multiple_of: int = 256,
        norm_eps: float = 1e-5,
        rope_theta: float = 10000.0,
        init_base_std: Optional[float] = 0.02,
        init_std_factor: str = "disabled",
        max_seqlen: int = 1024,
        rope_dim: int = 1,          # 0=NoPE, 1=1D-RoPE, 4=4D-RoPE
        tok_idx_type: Optional[str] = "t_coarse",
        # ── DecoderArgs ─────────────────────────────────────────────────────
        t_dim: int = 64,
        seqlen_t: bool = False,
        # ── DecoderTransformerArgs ───────────────────────────────────────────
        seed: int = 42,
        weight_tying: bool = False,
        sliding_window: int = 128,
        xattn_sliding_window: int = 32,
        input_dim: int = 64,
        decoder_encoder_dropout: float = 0.1,
        decoder_timestep_dropout: float = 0.1,
        encoder_sliding_window: int = 128,
        encoder_input_dim: Optional[int] = None,   # defaults to input_dim at runtime
        encoder_output_dim: Optional[int] = None,  # defaults to input_dim*2 at runtime
        encoder_latent_downsample_factor: int = 2,
        encoder_hidden_dim: Optional[int] = None,
        adaptive_loss_weighting: bool = False,
        num_fine_time_pts: int = 128,
        dont_noise_chan_xyz: bool = False,
        stft_global_sigma: Union[str, float] = 1.0,
        dropout_type: str = "zero",          # {"zero", "rand", "learnable"}
        bottleneck_type: str = "mmd",
        distill_output_dim: int = 0,
        codebook_size: int = 1024,
        levels: Optional[List[int]] = None,
        learnable_bias: bool = False,
        huber_c: Optional[float] = None,
        decoder_repa_index: float = float("inf"),
        encoder_repa_index: float = float("inf"),
        repa_dim: int = 1024,
        repa_loss_fn: str = "cosine",
        compression_free_conv_stem: bool = False,
        **kwargs,
    ):
        # ── BaseTransformerArgs ──────────────────────────────────────────────
        self.dim = dim
        self.n_layers = n_layers
        self.head_dim = head_dim
        self.n_heads = n_heads
        self.n_kv_heads = n_kv_heads
        self.ffn_dim_multiplier = ffn_dim_multiplier
        self.multiple_of = multiple_of
        self.norm_eps = norm_eps
        self.rope_theta = rope_theta
        self.init_base_std = init_base_std
        self.init_std_factor = init_std_factor
        self.max_seqlen = max_seqlen
        self.rope_dim = rope_dim
        self.tok_idx_type = tok_idx_type
        # ── DecoderArgs ─────────────────────────────────────────────────────
        self.t_dim = t_dim
        self.seqlen_t = seqlen_t
        # ── DecoderTransformerArgs ───────────────────────────────────────────
        self.seed = seed
        self.weight_tying = weight_tying
        self.sliding_window = sliding_window
        self.xattn_sliding_window = xattn_sliding_window
        self.input_dim = input_dim
        self.decoder_encoder_dropout = decoder_encoder_dropout
        self.decoder_timestep_dropout = decoder_timestep_dropout
        self.encoder_sliding_window = encoder_sliding_window
        self.encoder_input_dim = encoder_input_dim if encoder_input_dim is not None else input_dim
        self.encoder_output_dim = encoder_output_dim if encoder_output_dim is not None else input_dim * 2
        self.encoder_latent_downsample_factor = encoder_latent_downsample_factor
        self.encoder_hidden_dim = encoder_hidden_dim
        self.adaptive_loss_weighting = adaptive_loss_weighting
        self.num_fine_time_pts = num_fine_time_pts
        self.dont_noise_chan_xyz = dont_noise_chan_xyz
        self.stft_global_sigma = stft_global_sigma
        self.dropout_type = dropout_type
        self.bottleneck_type = bottleneck_type
        self.distill_output_dim = distill_output_dim
        self.codebook_size = codebook_size
        self.levels = levels if levels is not None else []
        self.learnable_bias = learnable_bias
        self.huber_c = huber_c
        self.decoder_repa_index = decoder_repa_index
        self.encoder_repa_index = encoder_repa_index
        self.repa_dim = repa_dim
        self.repa_loss_fn = repa_loss_fn
        self.compression_free_conv_stem = compression_free_conv_stem

        super().__init__(**kwargs)

    def to_decoder_transformer_args(self):
        """
        Convert back to a DecoderTransformerArgs dataclass instance so the raw
        Zyphra EncoderDecoder can be instantiated.
        """
        from .transformer import (
            DecoderTransformerArgs,
        )

        return DecoderTransformerArgs(
            dim=self.dim,
            n_layers=self.n_layers,
            head_dim=self.head_dim,
            n_heads=self.n_heads,
            n_kv_heads=self.n_kv_heads,
            ffn_dim_multiplier=self.ffn_dim_multiplier,
            multiple_of=self.multiple_of,
            norm_eps=self.norm_eps,
            rope_theta=self.rope_theta,
            init_base_std=self.init_base_std,
            init_std_factor=self.init_std_factor,
            max_seqlen=self.max_seqlen,
            rope_dim=self.rope_dim,
            tok_idx_type=self.tok_idx_type,
            t_dim=self.t_dim,
            seqlen_t=self.seqlen_t,
            seed=self.seed,
            weight_tying=self.weight_tying,
            sliding_window=self.sliding_window,
            xattn_sliding_window=self.xattn_sliding_window,
            input_dim=self.input_dim,
            decoder_encoder_dropout=self.decoder_encoder_dropout,
            decoder_timestep_dropout=self.decoder_timestep_dropout,
            encoder_sliding_window=self.encoder_sliding_window,
            encoder_input_dim=self.encoder_input_dim,
            encoder_output_dim=self.encoder_output_dim,
            encoder_latent_downsample_factor=self.encoder_latent_downsample_factor,
            encoder_hidden_dim=self.encoder_hidden_dim,
            adaptive_loss_weighting=self.adaptive_loss_weighting,
            num_fine_time_pts=self.num_fine_time_pts,
            dont_noise_chan_xyz=self.dont_noise_chan_xyz,
            stft_global_sigma=self.stft_global_sigma,
            dropout_type=self.dropout_type,
            bottleneck_type=self.bottleneck_type,
            distill_output_dim=self.distill_output_dim,
            codebook_size=self.codebook_size,
            levels=list(self.levels),
            learnable_bias=self.learnable_bias,
            huber_c=self.huber_c,
            decoder_repa_index=self.decoder_repa_index,
            encoder_repa_index=self.encoder_repa_index,
            repa_dim=self.repa_dim,
            repa_loss_fn=self.repa_loss_fn,
            compression_free_conv_stem=self.compression_free_conv_stem,
        )