""" ZunaModel: HuggingFace PreTrainedModel wrapper for the Zyphra ZUNA foundation model. Architecture (from Zyphra source): EncoderDecoder ├── encoder: EncoderTransformer → produces latent embeddings [B, L', encoder_output_dim] └── decoder: DecoderTransformer → flow-matching diffusion decoder The primary user-facing method is encode(), which runs only the encoder and returns the latent embeddings. forward() mirrors the Zyphra sample() signature for full encode→decode inference. Encoding call (verified against eeg_eval.py and EncoderDecoder.sample()): do_idx = (encoder_input.sum(axis=2) == 0).squeeze(0) enc_out, _ = model.model.encoder( token_values=encoder_input, seq_lens=seq_lens, tok_idx=tok_idx, do_idx=do_idx, ) # enc_out: [B, L/downsample_factor, encoder_output_dim] State-dict convention: the safetensors file ships keys prefixed with "model." which must be stripped before load_state_dict (see convert_weights.py). After stripping, keys match EncoderDecoder's submodule names directly: encoder.tok_embeddings.weight, decoder.layers.0.*, etc. """ from dataclasses import dataclass from typing import Optional import torch from transformers import PreTrainedModel from transformers.modeling_outputs import BaseModelOutput from .configuration_zuna import ZunaConfig class ZunaModel(PreTrainedModel): config_class = ZunaConfig # The raw Zyphra weights match EncoderDecoder's attribute tree directly # (after "model." prefix is stripped in convert_weights.py). base_model_prefix = "model" supports_gradient_checkpointing = False def __init__(self, config: ZunaConfig): super().__init__(config) from .transformer import ( EncoderDecoder, ) args = config.to_decoder_transformer_args() self.model = EncoderDecoder(args) # Store config fields needed at inference time. self.tok_idx_type = config.tok_idx_type self.rope_dim = config.rope_dim self.post_init() # ── internal helper ──────────────────────────────────────────────────── def _build_tok_idx( self, encoder_input: torch.Tensor, seq_lens: torch.Tensor, t_coarse: Optional[torch.Tensor], chan_id: Optional[torch.Tensor], chan_pos_discrete: Optional[torch.Tensor], ) -> Optional[torch.Tensor]: """ Replicate the tok_idx construction from EncoderDecoder.forward() (transformer.py lines 1001-1012). Only the types used by the public ZUNA checkpoint are implemented here; extend as needed. """ if self.tok_idx_type is None: return None elif self.tok_idx_type == "t_coarse" and self.rope_dim == 1: if t_coarse is None: raise ValueError("tok_idx_type='t_coarse' requires t_coarse tensor") return t_coarse elif self.tok_idx_type == "chan_id" and self.rope_dim == 1: if chan_id is None: raise ValueError("tok_idx_type='chan_id' requires chan_id tensor") return chan_id elif self.tok_idx_type == "stack_arange_seqlen" and self.rope_dim == 1: return torch.hstack( [torch.arange(sl) for sl in seq_lens] ).unsqueeze(0).unsqueeze(-1) elif self.tok_idx_type == "{x,y,z,tc}" and self.rope_dim == 4: if chan_pos_discrete is None or t_coarse is None: raise ValueError( "tok_idx_type='{x,y,z,tc}' requires chan_pos_discrete and t_coarse" ) return torch.cat((chan_pos_discrete, t_coarse), dim=2) else: raise ValueError( f"Unsupported tok_idx_type={self.tok_idx_type!r} with rope_dim={self.rope_dim}" ) # ── encode ──────────────────────────────────────────────────────────── @torch.no_grad() def encode( self, encoder_input: torch.Tensor, seq_lens: torch.Tensor, t_coarse: Optional[torch.Tensor] = None, chan_id: Optional[torch.Tensor] = None, chan_pos_discrete: Optional[torch.Tensor] = None, tok_idx: Optional[torch.Tensor] = None, ) -> torch.Tensor: """ Run only the encoder and return latent embeddings. Args: encoder_input: EEG signal tensor [B, seqlen, input_dim] seq_lens: Sequence lengths [B] (used for document masking) t_coarse: Coarse time index [B, seqlen, 1] (needed when tok_idx_type='t_coarse') chan_id: Channel IDs [B, seqlen, 1] (needed when tok_idx_type='chan_id') chan_pos_discrete: Discrete 3D positions [B, seqlen, 3] (needed for 4D-RoPE) tok_idx: Pre-built token index tensor; if provided, skips the automatic _build_tok_idx() construction above. Returns: Latent embeddings [B, seqlen // downsample_factor, encoder_output_dim] """ if encoder_input.ndim == 2: encoder_input = encoder_input.unsqueeze(0) if tok_idx is None: tok_idx = self._build_tok_idx( encoder_input, seq_lens, t_coarse, chan_id, chan_pos_discrete ) # Identify dropped-out channels: columns that are all-zero # (mirrors EncoderDecoder.sample() line 1055) do_idx = (encoder_input.sum(axis=2) == 0).squeeze(0) enc_out, _ = self.model.encoder( token_values=encoder_input, seq_lens=seq_lens, tok_idx=tok_idx, do_idx=do_idx, ) return enc_out # [B, L', encoder_output_dim] # ── forward ─────────────────────────────────────────────────────────── def forward( self, encoder_input: torch.Tensor, seq_lens: torch.Tensor, t_coarse: Optional[torch.Tensor] = None, chan_id: Optional[torch.Tensor] = None, chan_pos_discrete: Optional[torch.Tensor] = None, tok_idx: Optional[torch.Tensor] = None, sample_steps: int = 50, cfg: float = 1.0, return_latents: bool = False, ) -> BaseModelOutput: """ Full encode → diffusion-decode pass, mirroring EncoderDecoder.sample(). Args: encoder_input: EEG signal tensor [B, seqlen, input_dim] seq_lens: Sequence lengths [B] t_coarse: Coarse time index [B, seqlen, 1] chan_id: Channel IDs [B, seqlen, 1] chan_pos_discrete: Discrete 3D positions [B, seqlen, 3] tok_idx: Pre-built token index (overrides auto construction) sample_steps: Number of flow-matching diffusion steps (default 50) cfg: Classifier-free guidance scale (1.0 = no CFG) return_latents: If True, also return encoder latents in hidden_states Returns: BaseModelOutput with: last_hidden_state: reconstructed signal [B, seqlen, input_dim] hidden_states: (enc_out,) if return_latents else None """ if encoder_input.ndim == 2: encoder_input = encoder_input.unsqueeze(0) if tok_idx is None: tok_idx = self._build_tok_idx( encoder_input, seq_lens, t_coarse, chan_id, chan_pos_discrete ) reconstruction, _ = self.model.sample( encoder_input=encoder_input, seq_lens=seq_lens, tok_idx=tok_idx, sample_steps=sample_steps, cfg=cfg, ) hidden_states = None if return_latents: latents = self.encode( encoder_input=encoder_input, seq_lens=seq_lens, tok_idx=tok_idx, ) hidden_states = (latents,) return BaseModelOutput( last_hidden_state=reconstruction, hidden_states=hidden_states, )