modeling_zuna.py

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