add modeling_zeus.py

modeling_zeus.py

@@ -0,0 +1,759 @@
+import numpy as np
+import torch
+from torch import nn
+import math
+from typing import Tuple, Optional
+from transformers import PreTrainedModel
+from .configuration_zeus import ZeusConfig
+from basicts.modules import ACT2FN
+from basicts.modules.transformer import DecoderOnlyLayer, MultiHeadAttention, RotaryPositionEmbedding, AutoRegressiveDecoder
+from basicts.modules.norm import RMSNorm
+from flash_attn import flash_attn_varlen_func
+from flash_attn.bert_padding import unpad_input, pad_input
+
+
+class ZeusFlashAttention(nn.Module):
+    """
+    Encoder-only (BERT-style) Multi-Head Attention with FlashAttention v2
+    """
+    def __init__(
+        self,
+        hidden_size: int,
+        n_heads: int,
+        dropout: float = 0.0,
+        kv_heads: Optional[int] = None,
+        rope: Optional[torch.nn.Module] = None,
+    ):
+        super().__init__()
+        assert hidden_size % n_heads == 0
+
+        self.hidden_size = hidden_size
+        self.n_heads = n_heads
+        self.head_size = hidden_size // n_heads
+
+        self.q_proj = nn.Linear(hidden_size, hidden_size)
+        self.k_proj = nn.Linear(hidden_size, hidden_size)
+        self.v_proj = nn.Linear(hidden_size, hidden_size)
+        self.out_proj = nn.Linear(hidden_size, hidden_size, bias=False)
+
+        self.dropout_p = dropout
+        self.rope = rope
+
+    def _shape(self, x: torch.Tensor, B: int, L: int) -> torch.Tensor:
+        # [B, L, H*D] -> [B, L, H, D]
+        return x.view(B, L, self.n_heads, self.head_size)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_value: Optional[object] = None,
+        use_cache: bool = False,
+        output_attentions: bool = False,
+        layer_idx: Optional[int] = None,
+    ):
+        assert not output_attentions, \
+            "FlashAttention v2 does not support returning attention weights efficiently."
+
+        B, L, _ = hidden_states.shape
+        device = hidden_states.device
+
+        q = self._shape(self.q_proj(hidden_states), B, L)
+        k = self._shape(self.k_proj(hidden_states), B, L)
+        v = self._shape(self.v_proj(hidden_states), B, L)
+
+        if attention_mask is None:
+            mask = torch.ones((B, L), device=device, dtype=torch.bool)
+
+        q_unpad, indices, cu_seqlens, max_seqlen, _ = unpad_input(q, attention_mask)
+        k_unpad, _, _, _, _ = unpad_input(k, attention_mask)
+        v_unpad, _, _, _, _ = unpad_input(v, attention_mask)
+
+        if self.rope is not None:
+            if position_ids is None:
+                position_ids = torch.arange(L, device=device).unsqueeze(0).expand(B, -1)
+            pos = position_ids.reshape(-1)[indices]
+            q_unpad, k_unpad = self.rope(q_unpad, k_unpad, pos)
+
+        dropout_p = self.dropout_p if self.training else 0.0
+
+        attn_unpad = flash_attn_varlen_func(
+            q_unpad,
+            k_unpad,
+            v_unpad,
+            cu_seqlens_q=cu_seqlens,
+            cu_seqlens_k=cu_seqlens,
+            max_seqlen_q=max_seqlen,
+            max_seqlen_k=max_seqlen,
+            dropout_p=dropout_p,
+            causal=False,
+        )
+
+        attn_unpad = attn_unpad.reshape(-1, self.hidden_size)
+        context = pad_input(attn_unpad, indices, B, L)
+
+        output = self.out_proj(context)
+
+        return output, None, past_key_value
+
+
+class ZeusMLP(nn.Module):
+
+    def __init__(self, hidden_size: int, intermediate_size: int, hidden_act: str):
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
+        self.act_fn = ACT2FN[hidden_act]
+
+    def forward(self, hidden_state):
+        return self.down_proj(self.act_fn(self.gate_proj(hidden_state)) * self.up_proj(hidden_state))
+
+
+class ZeusInputEmbedding(nn.Module):
+
+    def __init__(self, input_size: int, hidden_size: int, hidden_act: str = "gelu"):
+        super().__init__()
+        self.input_size = input_size
+        self.hidden_size = hidden_size
+        self.intermediate_size = 4 * self.hidden_size
+        self.res_proj = nn.Linear(self.input_size, self.hidden_size, bias=False)
+        self.gate_proj = nn.Linear(self.input_size, self.intermediate_size, bias=True)
+        self.up_proj = nn.Linear(self.input_size, self.intermediate_size, bias=True)
+        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
+        self.act_fn = ACT2FN[hidden_act]
+
+    def forward(self, x: torch.Tensor):
+        return self.res_proj(x) + \
+            self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+
+
+class EncoderLayer(DecoderOnlyLayer):
+    def __init__(self, config: ZeusConfig, stage: int):
+        
+        attn_cls = ZeusFlashAttention \
+            if config.attn_implementation == "flash_attention_2" else MultiHeadAttention
+        
+        self_attn = attn_cls(
+            hidden_size=config.hidden_size[stage],
+            n_heads=config.n_heads[stage],
+            dropout=config.dropout,
+            rope=RotaryPositionEmbedding(
+                dim=config.hidden_size[stage] // config.n_heads[stage],
+                max_position_embeddings=4096
+            )
+        )
+        ffn_layer = ZeusMLP(
+            config.hidden_size[stage],
+            config.intermediate_size[stage],
+            config.hidden_act
+        )
+        super().__init__(self_attn, ffn_layer, (RMSNorm, config.hidden_size[stage]))
+
+
+class ZeusEncoder(AutoRegressiveDecoder):
+    def __init__(self, config: ZeusConfig, stage: int):
+        
+        decoder_layers = nn.ModuleList(
+            [
+                EncoderLayer(config, stage)
+                for _ in range(config.num_layers[stage])
+            ]  
+        )
+
+        layer_norm = RMSNorm(config.hidden_size[stage])
+        super().__init__(decoder_layers, layer_norm)
+
+        self.num_reg_tokens = config.num_reg_tokens
+
+        if self.num_reg_tokens > 0:
+            self.reg_tokens = nn.Parameter(
+                torch.randn(
+                    1, self.num_reg_tokens, config.hidden_size[stage]
+                ) * config.initializer_range
+            )
+    
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        **kwargs
+        ):
+
+        B, L, _ = hidden_states.size()
+        position_ids = torch.arange(
+            L,
+            dtype=torch.long,
+            device=hidden_states.device
+        ).unsqueeze(0)
+
+        if self.num_reg_tokens > 0:
+            reg_tokens = self.reg_tokens.expand(B, -1, -1)
+            hidden_states = torch.cat(
+                [reg_tokens, hidden_states], dim=1
+            )
+            position_ids = torch.cat(
+                [torch.zeros(
+                    1, self.num_reg_tokens,
+                    dtype=torch.long,
+                    device=hidden_states.device
+                ), position_ids], dim=1
+            )
+
+        hidden_states, attn_weights, kv_cache = super().forward(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            position_ids=position_ids.expand(B, -1),
+            **kwargs
+            )
+        
+        reg_tokens = hidden_states[:, :self.num_reg_tokens]
+        hidden_states = hidden_states[:, self.num_reg_tokens:]
+        
+        return hidden_states, attn_weights, kv_cache, reg_tokens
+
+class ZeusPoolingLayer(nn.Module):
+
+    def __init__(self, config: ZeusConfig, stage: int):
+        super().__init__()
+        self.stage = stage
+        self.config = config
+        self.factor = config.scales[stage] // config.scales[stage - 1]
+        self.proj = nn.Linear(
+            self.factor * config.hidden_size[stage - 1],
+            config.hidden_size[stage],
+            bias=False
+        )
+    
+    def forward(self, hidden_states: torch.Tensor, padding_mask: torch.Tensor):
+        batch_size, _, hidden_size = hidden_states.size()
+        hidden_states = hidden_states.reshape(batch_size, -1, self.factor * hidden_size)
+        hidden_states = self.proj(hidden_states)
+        padding_mask = padding_mask.reshape(batch_size, -1, self.factor, 1).any(dim=2)
+        return hidden_states, padding_mask
+
+class ZeusUnpoolingLayer(nn.Module):
+
+    def __init__(self, config: ZeusConfig, stage: int):
+        super().__init__()
+        self.stage = stage
+        self.config = config
+        self.factor = config.scales[stage - 1] // config.scales[stage]
+        self.proj = nn.Linear(
+            config.hidden_size[stage - 1],
+            self.factor * config.hidden_size[stage],
+            bias=False
+        )
+    
+    def forward(self, hidden_states: torch.Tensor, skip_connection: torch.Tensor):
+        batch_size, _, hidden_size = skip_connection.size()
+        hidden_states = self.proj(hidden_states)
+        hidden_states = hidden_states.reshape(batch_size, -1, hidden_size)
+        hidden_states = hidden_states + skip_connection
+        return hidden_states
+
+class ZeusPreTrainedModel(PreTrainedModel):
+    config_class = ZeusConfig
+
+    def _init_weights(self, module):
+        std = self.config.initializer_range
+        if isinstance(module, torch.nn.Linear):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, torch.nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+
+
+class Zeus(ZeusPreTrainedModel):
+
+    _supports_flash_attn_2 = True
+    
+    def __init__(self, config: ZeusConfig):
+        super().__init__(config)
+        self.config = config
+        self.scales = config.scales
+        self.num_reg_tokens = config.num_reg_tokens
+        self.num_scales = len(self.scales)
+
+        self.input_mlp = ZeusInputEmbedding(
+            config.input_dim,
+            config.hidden_size[0],
+            config.hidden_act
+            )
+
+        self.special_tokens = nn.Embedding(2, config.hidden_size[0])
+        self.pad_token_id = 0
+        self.mask_token_id = 1
+
+        self.encoders = nn.ModuleList()
+        self.downsamplers = nn.ModuleList()
+        self.upsamplers = nn.ModuleList()
+
+        # first layer
+        self.encoders.append(ZeusEncoder(config, 0))
+
+        # down samplers
+        for i in range(1, self.num_scales // 2 + 1):
+            self.encoders.append(ZeusEncoder(config, i))
+            self.downsamplers.append(ZeusPoolingLayer(config, i))
+        
+        for i in range(self.num_scales // 2 + 1, self.num_scales):
+            self.encoders.append(ZeusEncoder(config, i))
+            self.upsamplers.append(ZeusUnpoolingLayer(config, i))
+
+        self.num_quantiles = len(config.quantiles)
+        quantiles = torch.tensor(config.quantiles)
+        self.register_buffer("quantiles", quantiles, persistent=False)
+        self.head = nn.Linear(config.hidden_size[-1], self.num_quantiles)
+
+        self.post_init()
+
+    def _prepare_embedding(
+            self,
+            inputs: torch.Tensor,
+            targets_mask: torch.Tensor,
+            padding_mask: torch.Tensor = None,
+        ):
+
+        B, L, _ = inputs.shape
+        input_embeds = self.input_mlp(inputs) # [B, L, D]
+        
+        is_target = targets_mask == 1
+        input_embeds = torch.where(
+            is_target,
+            self.special_tokens(
+                torch.full_like(targets_mask.squeeze(-1), self.mask_token_id)
+            ),
+            input_embeds)
+        
+        if padding_mask is not None:
+            is_padding = padding_mask == 0
+            input_embeds = torch.where(
+                is_padding,
+                self.special_tokens(
+                    torch.full_like(padding_mask.squeeze(-1), self.pad_token_id)
+                ),
+                input_embeds)
+        if padding_mask is None:
+            padding_mask = torch.ones(
+                (B, L, 1), device=input_embeds.device, dtype=torch.long)
+        
+        # pad
+        max_scale = max(self.scales)
+        pad_len = math.ceil(L / max_scale) * max_scale - L
+        if pad_len > 0:
+            pad_tokens = self.special_tokens(
+                torch.full(
+                    (B, pad_len),
+                    self.pad_token_id,
+                    device=input_embeds.device
+                )
+            )
+            
+            input_embeds = torch.cat(
+                [input_embeds, pad_tokens],dim=1)
+            
+            padding_mask = torch.cat(
+                [
+                    padding_mask,
+                    torch.zeros(
+                        (B, pad_len, 1),
+                        device=input_embeds.device,
+                        dtype=padding_mask.dtype)
+                ],
+                dim=1
+            )
+            
+        return input_embeds, padding_mask
+    
+    def _prepare_attn_mask(
+        self,
+        hidden_states: torch.Tensor,
+        padding_mask: torch.Tensor = None,
+    ):
+        device = hidden_states.device
+        B, L, _ = hidden_states.shape
+        
+        if padding_mask is None:
+            padding_mask = torch.ones(
+                (B, L, 1), device=device, dtype=torch.long)
+
+        # reg tokens
+        if self.num_reg_tokens > 0:
+            attention_mask = torch.cat(
+                [
+                    torch.ones(
+                        (B, self.num_reg_tokens, 1),
+                        device=device,
+                        dtype=padding_mask.dtype
+                    ),
+                    padding_mask
+                ],
+                dim=1
+            )
+        else:
+            attention_mask = padding_mask
+        
+        if self.config.attn_implementation == "eager":
+            attention_mask = attention_mask.view(B, 1, 1, -1) # [B, 1, 1, L]
+            attention_mask = (1 - attention_mask.float()) * torch.finfo(hidden_states.dtype).min
+        else:
+            attention_mask = attention_mask.squeeze(-1) # [B, L]
+        return attention_mask
+    
+    def forward(
+        self, 
+        inputs: torch.Tensor,
+        targets_mask: Optional[torch.Tensor],
+        targets: Optional[torch.Tensor] = None,
+        padding_mask: Optional[torch.Tensor] = None,
+        return_all_hidden_states: bool = False
+    ):
+        """
+        x: [B, L, 1]
+        padding_mask: [B, L, 1] (0 for padding, 1 for valid)
+        target_mask: [B, L, 1] (1 for target/predict, 0 for context)
+        """
+
+        # embedding
+        ori_seq_len = inputs.shape[1]
+        ori_padding_mask = padding_mask
+        hidden_states, padding_mask = self._prepare_embedding(inputs, targets_mask, padding_mask)
+
+        scale_outputs = []
+        scale_padding_masks = []
+        all_hidden_states = []
+        reg_token_emb = None
+        
+        for i in range(self.num_scales):
+            
+            if i > 0:
+                
+                # pooling
+                if i <= self.num_scales // 2:
+                    scale_padding_masks.append(padding_mask)
+                    hidden_states, padding_mask = self.downsamplers[i - 1](hidden_states, padding_mask)
+            
+                # unpooling
+                else: # i > self.num_scales // 2
+                    idx = i - self.num_scales // 2 - 1
+                    hidden_states = self.upsamplers[idx](hidden_states, scale_outputs[self.num_scales - i - 1])
+                    padding_mask = scale_padding_masks[self.num_scales - i - 1]
+            
+            attention_mask = self._prepare_attn_mask(hidden_states, padding_mask)
+
+            hidden_states, _, _, reg_tokens = self.encoders[i](
+                hidden_states,
+                attention_mask=attention_mask
+                )
+            
+            if i == self.num_scales - 2:
+                reg_token_emb = reg_tokens.mean(dim=1) 
+            
+            if return_all_hidden_states:
+                all_hidden_states.append(hidden_states)
+
+            if i < self.num_scales:
+                scale_outputs.append(hidden_states)
+
+        # [B, L, D] -> [B, L, Q]
+        quantile_preds = self.head(hidden_states)[:, :ori_seq_len, :]
+
+        loss = 0.0
+        # target and not nan
+        if targets is not None:
+            loss_mask = (targets_mask * ori_padding_mask).float()
+            quantiles = self.quantiles.view(1, 1, self.num_quantiles).to(quantile_preds.dtype)
+            loss = 2 * torch.abs((targets - quantile_preds)
+                        * ((targets <= quantile_preds).float() - quantiles))
+            loss = loss * loss_mask
+            loss = loss.sum() / (loss_mask.sum() * self.num_quantiles)
+
+        return {
+            "prediction": quantile_preds,
+            "loss": loss,
+            "all_hidden_states": all_hidden_states,
+            "reg_token_emb": reg_token_emb,
+        }
+
+
+class ZeusForPrediction(Zeus):
+
+    def __init__(self, config: ZeusConfig):
+        super().__init__(config)
+
+    def generate(
+            self,
+            context: torch.Tensor,
+            prediction_length: int,
+            context_mask: torch.Tensor = None,
+            use_norm: bool = True
+            ) -> Tuple[torch.Tensor, torch.Tensor]:
+
+        context = context.to(self.device)
+        
+        ndim = context.ndim
+        num_features = None
+        if ndim == 2:
+            context = context.unsqueeze(-1)
+        elif ndim == 3 and context.shape[2] > 1:
+            _, L, num_features = context.shape
+            context = context.transpose(1, 2).view(-1, L, 1)
+        elif ndim == 1:
+            context = context.unsqueeze(0).unsqueeze(2)
+    
+        B, L, _ = context.shape
+        device = context.device
+
+        if use_norm:
+            mean = context.mean(dim=1, keepdim=True)
+            std = context.std(dim=1, keepdim=True)
+            context = (context - mean) / std
+            context = torch.arcsinh(context)
+        
+        inputs = torch.cat(
+            [context, torch.zeros(B, prediction_length, 1, device=device)], dim=1)
+        if context_mask is None:
+            context_mask = torch.torch.ones(B, L, 1, device=device, dtype=torch.int32)
+        padding_mask = torch.cat(
+            [
+                context_mask,
+                torch.ones(B, prediction_length, 1, dtype=torch.int32, device=device)
+            ], dim=1
+        )
+        targets_mask = torch.cat(
+            [
+                torch.zeros_like(context, dtype=torch.int32),
+                torch.ones(B, prediction_length, 1, dtype=torch.int32, device=device)
+            ], dim=1
+        )
+
+        with torch.autocast("cuda", dtype=torch.bfloat16):
+            outputs = self.forward(
+                inputs,
+                padding_mask=padding_mask,
+                targets_mask=targets_mask,
+            )
+
+        # [B, L, Q]
+        quantile_preds = outputs["prediction"][:, -prediction_length:, :]
+
+        if use_norm:
+            quantile_preds = torch.sinh(quantile_preds)
+            quantile_preds = quantile_preds * std + mean
+
+        # [B, L, 1]
+        prediction = quantile_preds.mean(dim=-1, keepdim=True)
+
+        if ndim == 2: # [B, L]
+            prediction = prediction.squeeze(-1)
+        elif ndim == 3 and num_features is not None:
+            # [B, L, N]
+            prediction = prediction.reshape(-1, num_features, prediction_length).transpose(1, 2)
+            prediction = quantile_preds.reshape(
+                -1, num_features, prediction_length, quantile_preds.shape[-1]
+                ).transpose(1, 2) # [B, L, N, Q]
+        elif ndim == 1:
+            prediction = prediction[0, :, 0] #[L,]
+            quantile_preds = quantile_preds[0] # [L, Q]
+
+        return prediction, quantile_preds
+    
+    def predict(
+            self,
+            context,
+            prediction_length,
+            use_norm: bool = True,
+            max_pred_len: int = 4096
+            ):
+
+        B = len(context)
+
+        series = []
+        Ns = []
+        for x in context:
+            if x.ndim == 1:          # [L] -> [1, L]
+                x = x[None, :]
+            else:                    # [L, N] -> [N, L]
+                x = x.T
+            series.append(x)
+            Ns.append(x.shape[0])
+
+        assert len(set(Ns)) == 1, "All arrays must have same N"
+        N = Ns[0]
+
+        padded = []
+        target_masks = []
+        for x in series:  # x: [N, L]
+            N_, L = x.shape
+            pad = np.full((N_, prediction_length), np.nan)
+            padded.append(np.concatenate([x, pad], axis=1))  # [N, L+F]
+
+            m = np.zeros((N_, L + prediction_length), dtype=bool)
+            m[:, L:] = 1
+            target_masks.append(m)
+
+        batch = []
+        # pad_masks = []
+        tgt_masks = []
+
+        for x, tm in zip(padded, target_masks):
+            N_, Lf = x.shape
+            if Lf >= max_pred_len:
+                x = x[:, -max_pred_len:]
+                tm = tm[:, -max_pred_len:]
+                # pm = np.ones((N_, max_pred_len), dtype=bool)
+            else:
+                pad_len = max_pred_len - Lf
+                x = np.concatenate([x, np.full((N_, pad_len), np.nan)], axis=1)
+                tm = np.concatenate([tm, np.zeros((N_, pad_len), bool)], axis=1)
+                # pm = np.concatenate([np.ones((N_, Lf)), np.zeros((N_, pad_len))], axis=1)
+
+            batch.append(x)
+            tgt_masks.append(tm)
+            # pad_masks.append(pm)
+
+        # [B, N, T] -> [B*N, T]
+        batch = np.stack(batch).reshape(B * N, max_pred_len, 1)
+        tgt_masks = np.stack(tgt_masks).reshape(B * N, max_pred_len, 1)
+        # pad_masks = np.stack(pad_masks).reshape(B * N, max_pred_len, 1)
+        pad_masks = (
+            (~np.isnan(batch))
+            | (tgt_masks.astype(bool))
+        ).astype(np.int32)
+
+        if use_norm:
+            mean = np.nanmean(batch, axis=1, keepdims=True)
+            std = np.nanstd(batch, axis=1, keepdims=True)
+            mean[np.isnan(mean)] = 0.0
+            std[np.isnan(std)] = 1.0
+            std[std < 1e-3] = 1.0
+            batch_norm = (batch - mean) / std
+            batch_norm = np.nan_to_num(batch_norm, nan=0.0)
+            batch_norm = np.arcsinh(batch_norm)
+        else:
+            batch_norm = np.nan_to_num(batch, nan=0.0)
+
+        x = torch.from_numpy(batch_norm).to(self.device).float()  # [B*N, T]
+        padding_mask = torch.from_numpy(pad_masks).int().to(self.device)
+        targets_mask = torch.from_numpy(tgt_masks).int().to(self.device)
+
+        # prediction: [B*N, T]
+        # quantile_prediction: [B*N, T, Q]
+        with torch.autocast("cuda", dtype=torch.bfloat16):
+            outputs = self.forward(
+                x,
+                padding_mask=padding_mask,
+                targets_mask=targets_mask,
+            )
+        
+        quantile_preds = outputs["prediction"].float().detach().cpu().numpy()  # [B*N, T, Q]
+        if use_norm:
+            quantile_preds = np.sinh(quantile_preds) * std + mean
+        quantile_preds = quantile_preds[tgt_masks.repeat(self.num_quantiles, axis=2)].reshape(B, N, prediction_length, self.num_quantiles)
+
+        preds = quantile_preds.mean(axis=-1)
+
+        if N == 1:
+            preds = preds[:, 0, :]
+            quantile_preds = quantile_preds[:, 0, :, :]
+
+        return preds, quantile_preds
+
+
+class ZeusForImputation(Zeus):
+    def __init__(self, config: ZeusConfig):
+        super().__init__(config)
+    
+    def generate(
+            self,
+            inputs: torch.Tensor,
+            targets_mask: torch.Tensor,
+            use_norm: bool = True
+            ) -> Tuple[torch.Tensor, torch.Tensor]:
+
+        # transform inputs and targets_mask to [B * N, L, 1]
+        ndim = inputs.ndim
+        num_features = None
+        if ndim == 2:
+            inputs = inputs.unsqueeze(-1)
+            targets_mask = targets_mask.unsqueeze(-1)
+        elif ndim == 3 and inputs.shape[2] > 1:
+            _, L, num_features = inputs.shape
+            inputs = inputs.transpose(1, 2).reshape(-1, L, 1)
+            targets_mask = targets_mask.transpose(1, 2).reshape(-1, L, 1)
+        elif ndim == 1:
+            inputs = inputs.unsqueeze(0).unsqueeze(2)
+            targets_mask = targets_mask.unsqueeze(0).unsqueeze(2)
+        
+        if use_norm:
+            inputs_mask = ~targets_mask # 1 for valid, 0 for invalid
+            valid_count = inputs_mask.sum(dim=1, keepdim=True).clamp_min(1)
+            mean = inputs.sum(dim=1, keepdim=True) / valid_count
+            inputs = (inputs - mean) * inputs_mask
+            std = torch.sqrt(
+                (inputs ** 2).sum(dim=1, keepdim=True) / valid_count + 1e-5)
+            inputs /= std
+            inputs = torch.arcsinh(inputs)
+
+        targets_mask = targets_mask.to(torch.int32)
+        with torch.autocast("cuda", dtype=torch.bfloat16):
+            outputs = self(inputs, targets_mask)
+        quantile_preds = outputs["prediction"]
+
+        if use_norm:
+            quantile_preds = torch.sinh(quantile_preds)
+            quantile_preds = quantile_preds * std + mean
+
+        if num_features is not None:
+            quantile_preds = quantile_preds.reshape(-1, num_features, L, self.config.quantiles).transpose(1, 2)
+
+        prediction = quantile_preds.mean(dim=-1, keepdim=True)
+        return prediction, quantile_preds
+
+
+class ZeusForClassification(Zeus):
+    def __init__(self, config: ZeusConfig):
+        super().__init__(config)
+    
+    def generate_one_sample(self, inputs: torch.Tensor, padding_mask: torch.Tensor = None, use_norm: bool = True):
+        # transform inputs and targets_mask to [B * N, L, 1]
+        B = inputs.shape[0]
+        ndim = inputs.ndim
+        num_features = None
+        if ndim == 2:
+            inputs = inputs.unsqueeze(-1)
+        elif ndim == 3 and inputs.shape[2] > 1:
+            _, L, num_features = inputs.shape
+            inputs = inputs.transpose(1, 2).view(-1, L, 1)
+        elif ndim == 1:
+            inputs = inputs.unsqueeze(0).unsqueeze(2)
+        
+        if use_norm:
+            if padding_mask is None:
+                padding_mask = torch.ones_like(inputs, dtype=torch.int32)
+            valid_count = padding_mask.sum(dim=1, keepdim=True).clamp_min(1)
+            mean = inputs.sum(dim=1, keepdim=True) / valid_count
+            inputs = (inputs - mean) * padding_mask
+            std = torch.sqrt(
+                (inputs ** 2).sum(dim=1, keepdim=True) / valid_count + 1e-5)
+            inputs /= std
+            inputs = torch.arcsinh(inputs)
+
+        targets_mask = torch.zeros_like(inputs, dtype=torch.int32)
+        with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
+            outputs = self(
+                inputs,
+                targets_mask=targets_mask,
+                padding_mask=padding_mask,
+                return_all_hidden_states=True
+            )
+        all_hidden_states = outputs["all_hidden_states"]
+
+        return all_hidden_states