""" 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, )