fastplms/models/dplm/modeling_dplm.py

"""FastPLMs-compatible DPLM implementation."""

# Copyright (c) 2024 Bytedance Ltd. and/or its affiliates
# SPDX-License-Identifier: Apache-2.0

from __future__ import annotations

import contextlib
import torch
import torch.nn as nn
from dataclasses import dataclass
from typing import ClassVar
from einops import rearrange
from torch.nn import functional as F
from torch.nn.attention import SDPBackend, sdpa_kernel
from transformers import EsmTokenizer
from transformers.modeling_outputs import (
    MaskedLMOutput,
    ModelOutput,
    SequenceClassifierOutput,
    TokenClassifierOutput,
)
from transformers.models.esm.configuration_esm import EsmConfig
from transformers.models.esm.modeling_esm import (
    EsmAttention,
    EsmClassificationHead,
    EsmContactPredictionHead,
    EsmEmbeddings,
    EsmEncoder,
    EsmLayer,
    EsmLMHead,
    EsmPooler,
    EsmPreTrainedModel,
    EsmSelfAttention,
)

from fastplms.models._diffusion_generation import generate_dplm
from fastplms.models._esm_rotary import RotaryEmbedding


try:
    from fastplms.attention import (
        AttentionBackend,
        BlockMask,
        FastPLMsAttentionMixin,
        _get_flex_attention_fn,
        flex_attention,
        get_attention_mask,
        kernels_flash_attention_func,
        resolve_attention_backend,
        resolve_attention_backend_for_call,
    )
    from fastplms.embeddings import EmbeddingMixin, select_hidden_state_embeddings
    from fastplms.models.ttt import FastPLMTestTimeTrainingMixin
except ModuleNotFoundError as error:
    _COMPOSITE_REQUIRED_NAMES = (
        "AttentionBackend",
        "BlockMask",
        "EmbeddingMixin",
        "FastPLMsAttentionMixin",
        "FastPLMTestTimeTrainingMixin",
        "_get_flex_attention_fn",
        "flex_attention",
        "get_attention_mask",
        "kernels_flash_attention_func",
        "resolve_attention_backend",
        "resolve_attention_backend_for_call",
        "select_hidden_state_embeddings",
    )
    if error.name != "fastplms" or any(
        name not in globals() for name in _COMPOSITE_REQUIRED_NAMES
    ):
        raise
    # Legacy flat Hub composites define every shared symbol above this block.


@dataclass
class DPLMMaskedLMOutput(MaskedLMOutput):
    """Masked-LM output with DPLM extensions after the HF fields."""

    s_max: tuple[list[torch.Tensor], ...] | None = None
    last_hidden_state: torch.Tensor | None = None


@dataclass
class DPLMSequenceClassifierOutput(SequenceClassifierOutput):
    """Sequence-classification output with optional attention diagnostics."""

    s_max: tuple[list[torch.Tensor], ...] | None = None


@dataclass
class DPLMTokenClassifierOutput(TokenClassifierOutput):
    """Token-classification output with optional attention diagnostics."""

    s_max: tuple[list[torch.Tensor], ...] | None = None


@dataclass
class DPLMEncoderOutput(ModelOutput):
    last_hidden_state: torch.Tensor | None = None
    pooler_output: torch.Tensor | None = None
    hidden_states: tuple[torch.Tensor, ...] | None = None
    attentions: tuple[torch.Tensor, ...] | None = None
    s_max: tuple[list[torch.Tensor], ...] | None = None


def _reject_unsupported_dplm_arguments(**arguments: object) -> None:
    unsupported = [
        name
        for name, value in arguments.items()
        if value is not None and not (name == "use_cache" and value is False)
    ]
    if unsupported:
        names = ", ".join(sorted(unsupported))
        raise ValueError(
            "DPLM is an encoder-only diffusion model and does not support "
            f"decoder, cross-attention, or KV-cache arguments: {names}."
        )


class DPLMConfig(EsmConfig):
    model_type = "dplm"

    def __init__(
        self,
        attn_backend: str | None = None,
        add_pooling_layer: bool = False,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.attn_backend = attn_backend
        self.add_pooling_layer = add_pooling_layer
        self.tie_word_embeddings = False


_TOKENIZER_LOAD_CONTEXT_KEYS = (
    "cache_dir",
    "force_download",
    "local_files_only",
    "proxies",
    "revision",
    "subfolder",
    "token",
    "trust_remote_code",
)


class DPLMPreTrainedModel(FastPLMsAttentionMixin, EsmPreTrainedModel):
    config_class = DPLMConfig
    # All advertised wrappers install the encoder at ``self.esm``.  Keep the
    # Hugging Face base-model and checkpoint-prefix contract aligned with that
    # actual module path.
    base_model_prefix = "esm"
    supports_gradient_checkpointing = True
    all_tied_weights_keys: ClassVar[dict[str, str]] = {}
    _supports_flash_attn = True
    _supports_flash_attn_2 = False
    _supports_flash_attn_3 = True
    _fastplms_attention_implementations = (
        "eager",
        "sdpa",
        "flex_attention",
        "flash_attention_3",
    )

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
        load_context = {key: kwargs[key] for key in _TOKENIZER_LOAD_CONTEXT_KEYS if key in kwargs}
        if "token" not in load_context and "use_auth_token" in kwargs:
            load_context["token"] = kwargs["use_auth_token"]
        load_context["source"] = pretrained_model_name_or_path

        loaded = super().from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs)
        model = loaded[0] if isinstance(loaded, tuple) else loaded
        model.__dict__["_fastplms_tokenizer_load_context"] = load_context
        model.__dict__["_fastplms_tokenizer"] = None
        return loaded

    @property
    def tokenizer(self):
        tokenizer = self.__dict__.get("_fastplms_tokenizer")
        if tokenizer is None:
            load_context = dict(self.__dict__.get("_fastplms_tokenizer_load_context") or {})
            source = load_context.pop("source", None)
            if source is None:
                source = str(getattr(self.config, "_name_or_path", "")).strip()
            if not source:
                raise RuntimeError(
                    "DPLM tokenizer loading requires a model loaded with from_pretrained "
                    "so checkpoint provenance is available."
                )
            tokenizer_kwargs = {
                key: value
                for key, value in load_context.items()
                if key in _TOKENIZER_LOAD_CONTEXT_KEYS and value is not None
            }
            resolved_revision = getattr(self.config, "_commit_hash", None)
            if resolved_revision:
                tokenizer_kwargs["revision"] = resolved_revision
            tokenizer = EsmTokenizer.from_pretrained(source, **tokenizer_kwargs)
            self.__dict__["_fastplms_tokenizer"] = tokenizer
        return tokenizer

    @tokenizer.setter
    def tokenizer(self, value) -> None:
        self.__dict__["_fastplms_tokenizer"] = value

    @property
    def attn_backend(self) -> str:
        return self.config.attn_backend

    @attn_backend.setter
    def attn_backend(self, backend: str) -> None:
        if backend not in self._fastplms_attention_implementations:
            raise ValueError(
                f"DPLM does not support {backend!r}; expected one of "
                f"{self._fastplms_attention_implementations}."
            )
        self.config.attn_backend = backend
        resolved = resolve_attention_backend(backend)
        for module in self.modules():
            if isinstance(module, ModifiedEsmEncoder):
                module.attention_backend = resolved
            elif isinstance(module, ModifiedEsmSelfAttention):
                module.attn_backend = resolved


class ModifiedEsmSelfAttention(EsmSelfAttention):
    def __init__(self, config, position_embedding_type=None) -> None:
        super().__init__(config, position_embedding_type)
        self.config = config
        self.scale = self.attention_head_size**-0.5
        self.dropout_prob = float(config.attention_probs_dropout_prob)
        self.attn_backend = resolve_attention_backend(config.attn_backend)
        if self.position_embedding_type == "rotary":
            self.rotary_embeddings = RotaryEmbedding(dim=self.attention_head_size)

    def transpose_for_scores(self, x: torch.Tensor) -> torch.Tensor:
        # x: (b, l, d)
        new_x_shape = (*x.size()[:-1], self.num_attention_heads, self.attention_head_size)
        x = x.view(new_x_shape)  # (b, l, h, d_h)
        return x.permute(0, 2, 1, 3)  # (b, h, l, d_h)

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask_2d: torch.Tensor | None = None,
        attention_mask_4d: torch.Tensor | None = None,
        flex_block_mask: object | None = None,
        head_mask: torch.FloatTensor | None = None,
        encoder_hidden_states: torch.FloatTensor | None = None,
        encoder_attention_mask: torch.FloatTensor | None = None,
        past_key_value: tuple[tuple[torch.FloatTensor]] | None = None,
        output_attentions: bool | None = False,
        output_s_max: bool | None = False,
        past_key_values: tuple[tuple[torch.FloatTensor]] | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor | None, list[torch.Tensor] | None]:
        # hidden_states: (b, l_q, d); encoder_hidden_states: (b, l_kv, d) or None
        if past_key_values is not None:
            past_key_value = past_key_values

        mixed_query_layer = self.query(hidden_states)  # (b, l_q, d)
        is_cross_attention = encoder_hidden_states is not None

        if is_cross_attention and past_key_value is not None:
            key_layer = past_key_value[0]  # (b, h, l_kv, d_h)
            value_layer = past_key_value[1]  # (b, h, l_kv, d_h)
            cross_attn_mask = encoder_attention_mask
        elif is_cross_attention:
            key_layer = self.transpose_for_scores(  # (b, h, l_kv, d_h)
                self.key(encoder_hidden_states)
            )
            value_layer = self.transpose_for_scores(  # (b, h, l_kv, d_h)
                self.value(encoder_hidden_states)
            )
            cross_attn_mask = encoder_attention_mask
        elif past_key_value is not None:
            key_layer = self.transpose_for_scores(self.key(hidden_states))  # (b, h, l_q, d_h)
            value_layer = self.transpose_for_scores(self.value(hidden_states))  # (b, h, l_q, d_h)
            key_layer = torch.cat([past_key_value[0], key_layer], dim=2)  # (b, h, l_kv, d_h)
            value_layer = torch.cat([past_key_value[1], value_layer], dim=2)  # (b, h, l_kv, d_h)
            cross_attn_mask = None
        else:
            key_layer = self.transpose_for_scores(self.key(hidden_states))  # (b, h, l_q, d_h)
            value_layer = self.transpose_for_scores(self.value(hidden_states))  # (b, h, l_q, d_h)
            cross_attn_mask = None

        query_layer = self.transpose_for_scores(mixed_query_layer) * self.scale  # (b, h, l_q, d_h)

        if self.position_embedding_type == "rotary":
            query_layer, key_layer = self.rotary_embeddings(  # Q: (b,h,l_q,d_h), K: (b,h,l_kv,d_h)
                query_layer,
                key_layer,
            )

        if self.position_embedding_type in ["relative_key", "relative_key_query"]:
            raise NotImplementedError

        query_layer = query_layer.contiguous()  # (b, h, l_q, d_h)
        key_layer = key_layer.contiguous()  # (b, h, l_kv, d_h)
        value_layer = value_layer.contiguous()  # (b, h, l_kv, d_h)

        if is_cross_attention:
            if self.attn_backend not in {
                AttentionBackend.EAGER,
                AttentionBackend.SDPA,
            }:
                raise RuntimeError(
                    f"DPLM cross-attention does not implement {self.attn_backend.value!r}. "
                    "Use eager or SDPA for decoder cross-attention."
                )
            if output_attentions:
                attn_output, attn_weights, s_max = self._manual_attn(
                    query_layer,
                    key_layer,
                    value_layer,
                    cross_attn_mask,
                    output_s_max,
                )
            elif self.attn_backend == AttentionBackend.EAGER:
                attn_output, _, s_max = self._manual_attn(
                    query_layer,
                    key_layer,
                    value_layer,
                    cross_attn_mask,
                    output_s_max,
                )
                attn_weights = None
            elif self.attn_backend == AttentionBackend.SDPA:
                attn_output, attn_weights = self._sdpa_attn(
                    query_layer,
                    key_layer,
                    value_layer,
                    cross_attn_mask,
                )
                s_max = self._compute_s_max(query_layer, key_layer) if output_s_max else None
        else:
            attn_output, attn_weights, s_max = self._attn(
                query_layer,
                key_layer,
                value_layer,
                attention_mask_2d=attention_mask_2d,
                attention_mask_4d=attention_mask_4d,
                flex_block_mask=flex_block_mask,
                output_attentions=output_attentions,
                output_s_max=output_s_max,
            )

        if head_mask is not None and torch.is_tensor(head_mask):
            batch_size, seq_len, _ = attn_output.shape
            attn_output = attn_output.view(  # (b, l_q, h, d_h)
                batch_size, seq_len, self.num_attention_heads, self.attention_head_size
            )
            attn_output = attn_output.permute(0, 2, 1, 3) * head_mask  # (b, h, l_q, d_h)
            attn_output = rearrange(attn_output, "b h s d -> b s (h d)")  # (b, l_q, d)

        return attn_output, attn_weights, s_max  # (b, l_q, d), optional weights, heads

    def _attn(
        self,
        query_heads: torch.Tensor,
        key_heads: torch.Tensor,
        value_heads: torch.Tensor,
        attention_mask_2d: torch.Tensor | None = None,
        attention_mask_4d: torch.Tensor | None = None,
        flex_block_mask: BlockMask | None = None,
        output_attentions: bool = False,
        output_s_max: bool = False,
    ) -> tuple[torch.Tensor, torch.Tensor | None, list[torch.Tensor] | None]:
        if output_attentions:
            return self._manual_attn(
                query_heads, key_heads, value_heads, attention_mask_4d, output_s_max
            )

        if (
            self.training
            and self.dropout_prob > 0
            and (self.attn_backend.is_flash or self.attn_backend == AttentionBackend.FLEX_ATTENTION)
        ):
            raise RuntimeError(
                f"DPLM {self.attn_backend.value} attention is inference-only when attention "
                "dropout is nonzero. Use eager or SDPA for this training configuration."
            )

        if self.attn_backend == AttentionBackend.EAGER:
            attn_output, _, s_max = self._manual_attn(
                query_heads, key_heads, value_heads, attention_mask_4d, output_s_max
            )
            return attn_output, None, s_max
        if self.attn_backend.is_flash:
            attn_output, attn_weights = self._kernels_flash_attn(
                query_heads, key_heads, value_heads, attention_mask_2d
            )
        elif self.attn_backend == AttentionBackend.FLEX:
            attn_output, attn_weights = self._flex_attn(
                query_heads,
                key_heads,
                value_heads,
                flex_block_mask,
                attention_mask_2d,
            )
        elif self.attn_backend == AttentionBackend.SDPA:
            attn_output, attn_weights = self._sdpa_attn(
                query_heads, key_heads, value_heads, attention_mask_4d
            )
        else:
            raise AssertionError(f"Unsupported resolved backend: {self.attn_backend}")

        s_max = self._compute_s_max(query_heads, key_heads) if output_s_max else None
        return attn_output, attn_weights, s_max

    @torch.no_grad()
    def _compute_s_max(
        self, query_heads: torch.Tensor, key_heads: torch.Tensor
    ) -> list[torch.Tensor]:
        q_norm = torch.linalg.vector_norm(query_heads, dim=-1)  # (b, h, l_q)
        k_norm = torch.linalg.vector_norm(key_heads, dim=-1)  # (b, h, l_kv)
        s_max_bound = (  # (h,)
            q_norm.max(dim=-1).values * k_norm.max(dim=-1).values
        ).max(dim=0).values
        return [s_max_bound[h] for h in range(self.num_attention_heads)]

    def _manual_attn(
        self,
        query_heads: torch.Tensor,
        key_heads: torch.Tensor,
        value_heads: torch.Tensor,
        attention_mask_4d: torch.Tensor | None = None,
        output_s_max: bool = False,
    ) -> tuple[torch.Tensor, torch.Tensor, list[torch.Tensor] | None]:
        attn_weights = torch.matmul(  # (b, h, l_q, l_kv)
            query_heads,
            key_heads.transpose(-1, -2),
        )
        if attention_mask_4d is not None:
            if attention_mask_4d.dtype == torch.bool:
                attn_weights = attn_weights.masked_fill(  # (b, h, l_q, l_kv)
                    attention_mask_4d.logical_not(),
                    float("-inf"),
                )
            else:
                attn_weights = attn_weights + attention_mask_4d.to(  # (b, h, l_q, l_kv)
                    device=attn_weights.device,
                    dtype=attn_weights.dtype,
                )
        attn_weights = F.softmax(attn_weights, dim=-1)  # (b, h, l_q, l_kv)
        if self.dropout_prob > 0 and self.training:
            attn_weights = F.dropout(  # (b, h, l_q, l_kv)
                attn_weights,
                p=self.dropout_prob,
                training=True,
            )
        context_heads = torch.matmul(attn_weights, value_heads)  # (b, h, l_q, d_h)
        attn_output = rearrange(context_heads, "b h s d -> b s (h d)")  # (b, l_q, d)
        s_max = self._compute_s_max(query_heads, key_heads) if output_s_max else None
        return attn_output, attn_weights, s_max

    def _kernels_flash_attn(
        self,
        query_heads: torch.Tensor,
        key_heads: torch.Tensor,
        value_heads: torch.Tensor,
        attention_mask_2d: torch.Tensor | None = None,
    ) -> tuple[torch.Tensor, None]:
        query_tokens = query_heads.transpose(1, 2).contiguous()  # (b, l_q, h, d_h)
        key_tokens = key_heads.transpose(1, 2).contiguous()  # (b, l_kv, h, d_h)
        value_tokens = value_heads.transpose(1, 2).contiguous()  # (b, l_kv, h, d_h)
        # Q has been pre-scaled by self.scale = 1/sqrt(head_dim) in forward().
        # Pass softmax_scale=1.0 to prevent double-scaling by the kernel.
        attn_output = kernels_flash_attention_func(  # (b, l_q, h, d_h)
            query_states=query_tokens,
            key_states=key_tokens,
            value_states=value_tokens,
            attention_mask_2d=attention_mask_2d,
            causal=False,
            softmax_scale=1.0,
            implementation=self.attn_backend.value,
        )
        return rearrange(attn_output, "b s h d -> b s (h d)"), None  # (b, l_q, d), None

    def _flex_attn(
        self,
        query_heads: torch.Tensor,
        key_heads: torch.Tensor,
        value_heads: torch.Tensor,
        flex_block_mask: BlockMask | None = None,
        attention_mask_2d: torch.Tensor | None = None,
    ) -> tuple[torch.Tensor, None]:
        if flex_attention is None:
            raise RuntimeError("Flex attention is not available in this environment.")
        fn = _get_flex_attention_fn(
            device=query_heads.device,
            dtype=query_heads.dtype,
            shape=tuple(query_heads.shape),
            mask_semantics="padding",
        )
        context_heads = fn(  # (b, h, l_q, d_h)
            query_heads, key_heads, value_heads, block_mask=flex_block_mask, scale=1.0
        )
        return rearrange(context_heads, "b h s d -> b s (h d)"), None  # (b, l_q, d), None

    def _sdpa_attn(
        self,
        query_heads: torch.Tensor,
        key_heads: torch.Tensor,
        value_heads: torch.Tensor,
        attention_mask_4d: torch.Tensor | None = None,
    ) -> tuple[torch.Tensor, None]:
        # The pinned official DPLM path uses Torch's efficient SDPA kernel for
        # its non-null padding mask. Torch 2.13 otherwise selects cuDNN on H100,
        # changing every downstream hidden state. Requiring the same public
        # SDPA kernel makes the official FP32-storage/BF16-autocast path exact.
        kernel_context = (
            sdpa_kernel(SDPBackend.EFFICIENT_ATTENTION)
            if query_heads.is_cuda
            else contextlib.nullcontext()
        )
        with kernel_context:
            context_heads = F.scaled_dot_product_attention(
                query_heads,
                key_heads,
                value_heads,
                attn_mask=attention_mask_4d,
                dropout_p=self.dropout_prob if self.training else 0.0,
                scale=1.0,
            )
        return rearrange(context_heads, "b h s d -> b s (h d)"), None


class ModifiedEsmAttention(EsmAttention):
    def __init__(self, config) -> None:
        # Reuse Transformers' maintained ESM container layout, replacing only
        # the self-attention engine that FastPLMs extends. This preserves the
        # checkpoint schema without duplicating an upstream DPLM constructor.
        EsmAttention.__init__(self, config)
        self.self = ModifiedEsmSelfAttention(config)

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask_2d: torch.Tensor | None = None,
        attention_mask_4d: torch.Tensor | None = None,
        flex_block_mask: object | None = None,
        head_mask: torch.Tensor | None = None,
        encoder_hidden_states: torch.Tensor | None = None,
        encoder_attention_mask: torch.Tensor | None = None,
        past_key_value: tuple[tuple[torch.FloatTensor]] | None = None,
        output_attentions: bool = False,
        output_s_max: bool = False,
    ) -> tuple[torch.Tensor, torch.Tensor | None, list[torch.Tensor] | None]:
        hidden_states_ln = self.LayerNorm(hidden_states)  # (b, l, d)
        attn_output, attn_weights, s_max = self.self(  # (b, l, d), optional weights, heads
            hidden_states_ln,
            attention_mask_2d=attention_mask_2d,
            attention_mask_4d=attention_mask_4d,
            flex_block_mask=flex_block_mask,
            head_mask=head_mask,
            encoder_hidden_states=encoder_hidden_states,
            encoder_attention_mask=encoder_attention_mask,
            past_key_value=past_key_value,
            output_attentions=output_attentions,
            output_s_max=output_s_max,
        )
        attention_output = self.output(attn_output, hidden_states)  # (b, l, d)
        return attention_output, attn_weights, s_max


class ModifiedEsmLayer(EsmLayer):
    def __init__(self, config) -> None:
        # Transformers owns the feed-forward, normalization, and decoder
        # plumbing. Only attention dispatch differs for DPLM.
        EsmLayer.__init__(self, config)
        self.attention = ModifiedEsmAttention(config)
        if self.add_cross_attention:
            self.crossattention = ModifiedEsmAttention(config)

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask_2d: torch.Tensor | None = None,
        attention_mask_4d: torch.Tensor | None = None,
        flex_block_mask: object | None = None,
        head_mask: torch.Tensor | None = None,
        encoder_hidden_states: torch.Tensor | None = None,
        encoder_attention_mask: torch.Tensor | None = None,
        past_key_value: tuple[tuple[torch.FloatTensor]] | None = None,
        output_attentions: bool = False,
        output_s_max: bool = False,
    ) -> tuple[torch.Tensor, torch.Tensor | None, list[torch.Tensor] | None]:
        attention_output, attn_weights, s_max = self.attention(  # (b, l, d), weights, heads
            hidden_states,
            attention_mask_2d=attention_mask_2d,
            attention_mask_4d=attention_mask_4d,
            flex_block_mask=flex_block_mask,
            head_mask=head_mask,
            output_attentions=output_attentions,
            output_s_max=output_s_max,
            past_key_value=past_key_value[:2] if past_key_value is not None else None,
        )

        if self.is_decoder and encoder_hidden_states is not None:
            if self.add_cross_attention is False:
                raise AttributeError(
                    f"If `encoder_hidden_states` are passed, {self} has to be "
                    "instantiated with cross-attention "
                    "layers by setting `config.add_cross_attention=True`"
                )
            cross_attn_past_key_value = past_key_value[-2:] if past_key_value is not None else None
            cross_attention_output, _, _ = self.crossattention(  # (b, l, d), weights, heads
                attention_output,
                attention_mask_2d=attention_mask_2d,
                attention_mask_4d=attention_mask_4d,
                head_mask=head_mask,
                encoder_hidden_states=encoder_hidden_states,
                encoder_attention_mask=encoder_attention_mask,
                past_key_value=cross_attn_past_key_value,
                output_attentions=output_attentions,
                output_s_max=False,
            )
            attention_output = cross_attention_output  # (b, l, d)

        layer_output = self.feed_forward_chunk(attention_output)  # (b, l, d)
        return layer_output, attn_weights, s_max


class ModifiedEsmEncoder(EsmEncoder):
    def __init__(self, config):
        # Start from the public Transformers encoder contract, then substitute
        # backend-aware layers while retaining every canonical state key.
        EsmEncoder.__init__(self, config)
        self.attention_backend = resolve_attention_backend(config.attn_backend)
        self.layer = nn.ModuleList(
            ModifiedEsmLayer(config) for _ in range(config.num_hidden_layers)
        )

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: torch.Tensor | None = None,
        head_mask: torch.Tensor | None = None,
        encoder_hidden_states: torch.Tensor | None = None,
        encoder_attention_mask: torch.Tensor | None = None,
        past_key_values: list[tuple[tuple[torch.FloatTensor]]] | None = None,
        use_cache: bool | None = None,
        output_attentions: bool = False,
        output_hidden_states: bool = False,
        output_s_max: bool = False,
    ) -> DPLMEncoderOutput:
        first_parameter = next(self.parameters(), None)
        if (
            not self.training
            and first_parameter is not None
            and first_parameter.dtype == torch.bfloat16
        ):
            raise RuntimeError(
                "DPLM BF16 inference requires FP32-resident parameters under "
                "CUDA BF16 autocast; static BF16 parameters do not meet the "
                "declared parity contract."
            )
        all_hidden_states = () if output_hidden_states else None
        all_self_attentions = () if output_attentions else None
        full_s_max = () if output_s_max else None

        effective_backend = resolve_attention_backend_for_call(
            self.attention_backend,
            output_attentions=output_attentions,
        )
        attention_mask_2d, attention_mask_4d, flex_block_mask = get_attention_mask(
            effective_backend=effective_backend,
            batch_size=hidden_states.shape[0],
            seq_len=hidden_states.shape[1],
            device=hidden_states.device,
            attention_mask=attention_mask,
            dtype=hidden_states.dtype,
            mask_semantics="padding",
        )

        for i, layer_module in enumerate(self.layer):
            if output_hidden_states:
                all_hidden_states = (*all_hidden_states, hidden_states)

            layer_head_mask = head_mask[i] if head_mask is not None else None
            past_key_value = past_key_values[i] if past_key_values is not None else None

            if self.gradient_checkpointing and self.training:
                hidden_states, attn_weights, s_max = self._gradient_checkpointing_func(
                    layer_module.__call__,
                    hidden_states,
                    attention_mask_2d,
                    attention_mask_4d,
                    flex_block_mask,
                    layer_head_mask,
                    encoder_hidden_states,
                    encoder_attention_mask,
                    past_key_value,
                    output_attentions,
                    output_s_max,
                )
            else:
                hidden_states, attn_weights, s_max = layer_module(
                    hidden_states,
                    attention_mask_2d=attention_mask_2d,
                    attention_mask_4d=attention_mask_4d,
                    flex_block_mask=flex_block_mask,
                    head_mask=layer_head_mask,
                    encoder_hidden_states=encoder_hidden_states,
                    encoder_attention_mask=encoder_attention_mask,
                    past_key_value=past_key_value,
                    output_attentions=output_attentions,
                    output_s_max=output_s_max,
                )

            if all_self_attentions is not None:
                all_self_attentions = (*all_self_attentions, attn_weights)
            if full_s_max is not None:
                full_s_max = (*full_s_max, s_max)

        if self.emb_layer_norm_after:
            hidden_states = self.emb_layer_norm_after(hidden_states)

        if output_hidden_states:
            all_hidden_states = (*all_hidden_states, hidden_states)

        return DPLMEncoderOutput(
            last_hidden_state=hidden_states,
            hidden_states=all_hidden_states,
            attentions=all_self_attentions,
            s_max=full_s_max,
        )


class FAST_DPLM_ENCODER(DPLMPreTrainedModel, EmbeddingMixin):
    """Inner encoder class that holds the actual ESM-style weights (embeddings, encoder,
    contact_head) so that the weight keys are prefixed with 'esm.' in the outer DPLMModel,
    matching pretrained DPLM checkpoints."""

    def __init__(self, config, **kwargs):
        DPLMPreTrainedModel.__init__(self, config, **kwargs)
        self.config = config
        self.embeddings = EsmEmbeddings(config)
        self.encoder = ModifiedEsmEncoder(config)
        self.contact_head = EsmContactPredictionHead(
            in_features=config.num_hidden_layers * config.num_attention_heads,
            bias=True,
        )
        self.post_init()

    def get_input_embeddings(self) -> nn.Module:
        return self.embeddings.word_embeddings

    def set_input_embeddings(self, value):
        self.embeddings.word_embeddings = value

    def _embed(
        self,
        input_ids: torch.Tensor,
        attention_mask: torch.Tensor | None = None,
        hidden_state_index: int = -1,
        store_all_hidden_states: bool = False,
    ) -> torch.Tensor:
        if attention_mask is None:
            attention_mask = input_ids.ne(self.config.pad_token_id)
        embedding_output = self.embeddings(input_ids, attention_mask=attention_mask)
        output_hidden_states = store_all_hidden_states or hidden_state_index != -1
        encoder_outputs = self.encoder(
            embedding_output,
            attention_mask=attention_mask,
            output_hidden_states=output_hidden_states,
            output_attentions=False,
        )
        return select_hidden_state_embeddings(
            encoder_outputs.last_hidden_state,
            encoder_outputs.hidden_states,
            hidden_state_index=hidden_state_index,
            store_all_hidden_states=store_all_hidden_states,
        )

    def predict_contacts(
        self, input_ids: torch.Tensor, attention_mask: torch.Tensor
    ) -> torch.Tensor:
        attns = self(input_ids, attention_mask=attention_mask, output_attentions=True).attentions
        attns = torch.stack(attns, dim=1)
        attns *= attention_mask.unsqueeze(1).unsqueeze(2).unsqueeze(3)
        attns *= attention_mask.unsqueeze(1).unsqueeze(2).unsqueeze(4)
        return self.contact_head(input_ids, attns)

    def _convert_head_mask_to_5d(
        self, head_mask: torch.Tensor, num_hidden_layers: int
    ) -> torch.Tensor:
        if head_mask.dim() == 1:
            head_mask = head_mask.unsqueeze(0).unsqueeze(0).unsqueeze(-1).unsqueeze(-1)
            head_mask = head_mask.expand(num_hidden_layers, -1, -1, -1, -1)
        elif head_mask.dim() == 2:
            head_mask = head_mask.unsqueeze(1).unsqueeze(-1).unsqueeze(-1)
        if head_mask.dim() != 5:
            raise ValueError(f"head_mask.dim != 5, got {head_mask.dim()}")
        head_mask = head_mask.to(dtype=self.dtype)
        return head_mask

    def get_head_mask(
        self,
        head_mask: torch.Tensor | None,
        num_hidden_layers: int,
        is_attention_chunked: bool = False,
    ) -> torch.Tensor | list[None]:
        if head_mask is None:
            return [None] * num_hidden_layers
        head_mask = self._convert_head_mask_to_5d(head_mask, num_hidden_layers)
        if is_attention_chunked:
            head_mask = head_mask.unsqueeze(-1)
        return head_mask

    def forward(
        self,
        input_ids: torch.Tensor | None = None,
        attention_mask: torch.Tensor | None = None,
        position_ids: torch.Tensor | None = None,
        head_mask: torch.Tensor | None = None,
        inputs_embeds: torch.Tensor | None = None,
        encoder_hidden_states: torch.Tensor | None = None,
        encoder_attention_mask: torch.Tensor | None = None,
        past_key_values: list[torch.FloatTensor] | None = None,
        use_cache: bool | None = None,
        output_attentions: bool | None = None,
        output_hidden_states: bool | None = None,
        output_s_max: bool | None = False,
        return_dict: bool | None = None,
    ) -> tuple[torch.Tensor] | DPLMEncoderOutput:
        if self.config.is_decoder or self.config.add_cross_attention:
            raise ValueError(
                "DPLM is encoder-only; is_decoder and add_cross_attention must be false."
            )
        _reject_unsupported_dplm_arguments(
            encoder_hidden_states=encoder_hidden_states,
            encoder_attention_mask=encoder_attention_mask,
            past_key_values=past_key_values,
            use_cache=use_cache,
        )
        output_attentions = (
            output_attentions if output_attentions is not None else self.config.output_attentions
        )
        output_hidden_states = (
            output_hidden_states
            if output_hidden_states is not None
            else self.config.output_hidden_states
        )
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        if self.config.is_decoder:
            use_cache = use_cache if use_cache is not None else self.config.use_cache
        else:
            use_cache = False

        if input_ids is not None and inputs_embeds is not None:
            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
        if input_ids is not None:
            input_shape = input_ids.size()
        elif inputs_embeds is not None:
            input_shape = inputs_embeds.size()[:-1]
        else:
            raise ValueError("You have to specify either input_ids or inputs_embeds")

        batch_size, seq_length = input_shape
        device = input_ids.device if input_ids is not None else inputs_embeds.device

        expected_attention_mask_shape = (batch_size, seq_length)
        if attention_mask is None:
            attention_mask_2d = torch.ones((batch_size, seq_length), device=device).bool()
        elif attention_mask.dim() == 4:
            raise ValueError(
                "DPLM accepts a two-dimensional padding mask. Passing a four-dimensional "
                "custom attention mask is unsupported because it cannot be applied to both "
                "the embedding and optimized-attention paths without changing semantics."
            )
        elif (
            attention_mask.dim() != 2
            or tuple(attention_mask.shape) != expected_attention_mask_shape
        ):
            raise ValueError(
                f"attention_mask must have shape {expected_attention_mask_shape}; "
                f"received {tuple(attention_mask.shape)}."
            )
        else:
            attention_mask_2d = attention_mask.to(device=device, dtype=torch.bool)

        encoder_extended_attention_mask = encoder_attention_mask
        if self.config.is_decoder and encoder_hidden_states is not None:
            encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states.size()
            encoder_hidden_shape = (encoder_batch_size, encoder_sequence_length)
            if encoder_attention_mask is None:
                encoder_attention_mask = torch.ones(encoder_hidden_shape, device=device)
            encoder_extended_attention_mask = self.invert_attention_mask(encoder_attention_mask)

        head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)

        embedding_output = self.embeddings(
            input_ids=input_ids,
            position_ids=position_ids,
            attention_mask=attention_mask_2d,
            inputs_embeds=inputs_embeds,
        )
        encoder_outputs = self.encoder(
            embedding_output,
            attention_mask=attention_mask_2d,
            head_mask=head_mask,
            encoder_hidden_states=encoder_hidden_states,
            encoder_attention_mask=encoder_extended_attention_mask,
            past_key_values=past_key_values,
            use_cache=use_cache,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            output_s_max=output_s_max,
        )
        sequence_output = encoder_outputs.last_hidden_state

        if return_dict is False:
            return (sequence_output, *encoder_outputs[1:])

        result = DPLMEncoderOutput(
            last_hidden_state=sequence_output,
            hidden_states=encoder_outputs.hidden_states,
            attentions=encoder_outputs.attentions,
            s_max=encoder_outputs.s_max,
        )
        return result


class DPLMModel(DPLMPreTrainedModel, EmbeddingMixin):
    config_class = DPLMConfig

    def __init__(self, config, add_pooling_layer: bool | None = None) -> None:
        DPLMPreTrainedModel.__init__(self, config)
        self.config = config
        self.esm = FAST_DPLM_ENCODER(config)
        if add_pooling_layer is None:
            add_pooling_layer = config.add_pooling_layer
        config.add_pooling_layer = bool(add_pooling_layer)
        self.pooler = EsmPooler(config) if add_pooling_layer else None
        self.post_init()

    def get_input_embeddings(self) -> nn.Module:
        return self.esm.embeddings.word_embeddings

    def set_input_embeddings(self, value):
        self.esm.embeddings.word_embeddings = value

    def _embed(
        self,
        input_ids: torch.Tensor,
        attention_mask: torch.Tensor | None = None,
        hidden_state_index: int = -1,
        store_all_hidden_states: bool = False,
    ) -> torch.Tensor:
        return self.esm._embed(
            input_ids,
            attention_mask,
            hidden_state_index=hidden_state_index,
            store_all_hidden_states=store_all_hidden_states,
        )

    def predict_contacts(
        self, input_ids: torch.Tensor, attention_mask: torch.Tensor
    ) -> torch.Tensor:
        return self.esm.predict_contacts(input_ids, attention_mask)

    def forward(
        self,
        input_ids: torch.Tensor | None = None,
        attention_mask: torch.Tensor | None = None,
        position_ids: torch.Tensor | None = None,
        head_mask: torch.Tensor | None = None,
        inputs_embeds: torch.Tensor | None = None,
        encoder_hidden_states: torch.Tensor | None = None,
        encoder_attention_mask: torch.Tensor | None = None,
        past_key_values: list[torch.FloatTensor] | None = None,
        use_cache: bool | None = None,
        output_attentions: bool | None = None,
        output_hidden_states: bool | None = None,
        output_s_max: bool | None = False,
        return_dict: bool | None = None,
    ) -> tuple[torch.Tensor] | DPLMEncoderOutput:
        outputs = self.esm(
            input_ids=input_ids,
            attention_mask=attention_mask,
            position_ids=position_ids,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
            encoder_hidden_states=encoder_hidden_states,
            encoder_attention_mask=encoder_attention_mask,
            past_key_values=past_key_values,
            use_cache=use_cache,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            output_s_max=output_s_max,
            return_dict=True,
        )
        sequence_output = outputs[0]  # (b, l, d)
        pooled_output = self.pooler(sequence_output) if self.pooler is not None else None  # (b, d)

        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
        result = DPLMEncoderOutput(
            last_hidden_state=sequence_output,
            pooler_output=pooled_output,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
            s_max=outputs.s_max,
        )
        return result if return_dict else result.to_tuple()


class DPLMForMaskedLM(FastPLMTestTimeTrainingMixin, DPLMPreTrainedModel, EmbeddingMixin):
    config_class = DPLMConfig

    def __init__(self, config, dropout: float | None = None) -> None:
        if dropout is not None:
            config.hidden_dropout_prob = dropout
        DPLMPreTrainedModel.__init__(self, config)
        self.esm = FAST_DPLM_ENCODER(config)
        self.lm_head = EsmLMHead(config)
        self.loss_fct = nn.CrossEntropyLoss()
        self.post_init()
        self.contact_head = None
        self.init_ttt({"lora_target_replace_module": "ModifiedEsmAttention"})

    def get_input_embeddings(self) -> nn.Module:
        return self.esm.get_input_embeddings()

    def set_input_embeddings(self, value: nn.Module) -> None:
        self.esm.set_input_embeddings(value)

    def get_output_embeddings(self):
        return self.lm_head.decoder

    def set_output_embeddings(self, new_embeddings):
        old_bias = self.lm_head.bias
        new_vocab_size = int(new_embeddings.out_features)
        if old_bias.shape[0] != new_vocab_size:
            resized_bias = old_bias.new_zeros(new_vocab_size)
            copy_length = min(old_bias.shape[0], new_vocab_size)
            with torch.no_grad():
                resized_bias[:copy_length].copy_(old_bias[:copy_length])
            self.lm_head.bias = nn.Parameter(resized_bias)
        # EsmLMHead.forward adds this standalone bias after the decoder. HF's
        # generic LM-head resizer may create a biased Linear, which would apply
        # the bias twice and introduce an undeclared shared tensor on save.
        new_embeddings.bias = None
        self.lm_head.decoder = new_embeddings

    def generate(
        self,
        input_tokens: torch.Tensor,
        tokenizer: object | None = None,
        max_iter: int | None = None,
        temperature: float | None = None,
        partial_masks: torch.Tensor | None = None,
        sampling_strategy: str = "gumbel_argmax",
        disable_resample: bool = False,
        resample_ratio: float = 0.25,
        show_progress: bool = False,
        **kwargs,
    ) -> torch.Tensor:
        """Generate protein tokens with the official DPLM diffusion schedule.

        ``input_tokens`` is X with shape (b, l). Positions marked ``True`` in
        ``partial_masks`` remain fixed. ``max_iter=None`` uses the official
        500-step schedule; shorter schedules are useful for rapid exploration.
        """

        if kwargs:
            names = ", ".join(sorted(kwargs))
            raise TypeError(f"Unexpected DPLM generation arguments: {names}")
        return generate_dplm(
            self,
            input_tokens,
            tokenizer=tokenizer,
            max_iter=max_iter,
            temperature=temperature,
            partial_masks=partial_masks,
            sampling_strategy=sampling_strategy,
            disable_resample=disable_resample,
            resample_ratio=resample_ratio,
            show_progress=show_progress,
        )

    def _embed(
        self,
        input_ids: torch.Tensor,
        attention_mask: torch.Tensor | None = None,
        hidden_state_index: int = -1,
        store_all_hidden_states: bool = False,
    ) -> torch.Tensor:
        return self.esm._embed(
            input_ids,
            attention_mask,
            hidden_state_index=hidden_state_index,
            store_all_hidden_states=store_all_hidden_states,
        )

    def predict_contacts(
        self, input_ids: torch.Tensor, attention_mask: torch.Tensor
    ) -> torch.Tensor:
        return self.esm.predict_contacts(input_ids, attention_mask=attention_mask)

    def _ttt_get_trainable_modules(self) -> list[nn.Module]:
        return [self.esm]

    def forward(
        self,
        input_ids: torch.Tensor | None = None,
        attention_mask: torch.Tensor | None = None,
        inputs_embeds: torch.Tensor | None = None,
        decoder_input_ids: torch.Tensor | None = None,
        decoder_attention_mask: torch.Tensor | None = None,
        decoder_inputs_embeds: torch.Tensor | None = None,
        labels: torch.Tensor | None = None,
        output_attentions: bool | None = None,
        output_hidden_states: bool | None = None,
        output_s_max: bool | None = False,
        return_dict: bool | None = None,
        encoder_hidden_states: torch.Tensor | None = None,
        encoder_attention_mask: torch.Tensor | None = None,
    ) -> tuple[torch.Tensor] | DPLMMaskedLMOutput:
        _reject_unsupported_dplm_arguments(
            decoder_input_ids=decoder_input_ids,
            decoder_attention_mask=decoder_attention_mask,
            decoder_inputs_embeds=decoder_inputs_embeds,
            encoder_hidden_states=encoder_hidden_states,
            encoder_attention_mask=encoder_attention_mask,
        )
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
        if attention_mask is None and input_ids is not None:
            attention_mask = input_ids.ne(self.config.pad_token_id)

        outputs = self.esm(
            input_ids=input_ids,
            attention_mask=attention_mask,
            inputs_embeds=inputs_embeds,
            encoder_hidden_states=encoder_hidden_states,
            encoder_attention_mask=encoder_attention_mask,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            output_s_max=output_s_max,
            return_dict=True,
        )
        sequence_output = outputs.last_hidden_state  # (b, l, d)
        logits = self.lm_head(sequence_output)  # (b, l, c)

        loss = None
        if labels is not None:
            labels = labels.to(logits.device)  # (b, l)
            loss = self.loss_fct(  # ()
                logits.view(-1, self.config.vocab_size),
                labels.view(-1),
            )

        result = DPLMMaskedLMOutput(
            loss=loss,
            logits=logits,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
            s_max=outputs.s_max,
            last_hidden_state=sequence_output,
        )
        return result if return_dict else result.to_tuple()


class DPLMForSequenceClassification(DPLMPreTrainedModel, EmbeddingMixin):
    config_class = DPLMConfig

    def get_input_embeddings(self) -> nn.Module:
        return self.esm.get_input_embeddings()

    def set_input_embeddings(self, value: nn.Module) -> None:
        self.esm.set_input_embeddings(value)

    def __init__(self, config) -> None:
        DPLMPreTrainedModel.__init__(self, config)
        self.num_labels = config.num_labels
        self.esm = FAST_DPLM_ENCODER(config)
        self.classifier = EsmClassificationHead(config)
        self.mse = nn.MSELoss()
        self.ce = nn.CrossEntropyLoss()
        self.bce = nn.BCEWithLogitsLoss()
        self.post_init()

    def _embed(
        self,
        input_ids: torch.Tensor,
        attention_mask: torch.Tensor | None = None,
        hidden_state_index: int = -1,
        store_all_hidden_states: bool = False,
    ) -> torch.Tensor:
        return self.esm._embed(
            input_ids,
            attention_mask,
            hidden_state_index=hidden_state_index,
            store_all_hidden_states=store_all_hidden_states,
        )

    def forward(
        self,
        input_ids: torch.Tensor | None = None,
        attention_mask: torch.Tensor | None = None,
        inputs_embeds: torch.Tensor | None = None,
        labels: torch.Tensor | None = None,
        output_attentions: bool | None = None,
        output_hidden_states: bool | None = None,
        output_s_max: bool | None = False,
        return_dict: bool | None = None,
    ) -> tuple[torch.Tensor, ...] | DPLMSequenceClassifierOutput:
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
        outputs = self.esm(
            input_ids=input_ids,
            attention_mask=attention_mask,
            inputs_embeds=inputs_embeds,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            output_s_max=output_s_max,
            return_dict=True,
        )
        sequence_output = outputs.last_hidden_state  # (b, l, d)
        logits = self.classifier(sequence_output)  # (b, c)

        loss = None
        if labels is not None:
            labels = labels.to(logits.device)  # (b,) or (b, c)
            if self.config.problem_type is None:
                if self.num_labels == 1:
                    self.config.problem_type = "regression"
                elif self.num_labels > 1 and (
                    labels.dtype == torch.long or labels.dtype == torch.int
                ):
                    self.config.problem_type = "single_label_classification"
                else:
                    self.config.problem_type = "multi_label_classification"

            if self.config.problem_type == "regression":
                if self.num_labels == 1:
                    loss = self.mse(logits.squeeze(), labels.squeeze())
                else:
                    loss = self.mse(logits, labels)
            elif self.config.problem_type == "single_label_classification":
                loss = self.ce(logits.view(-1, self.num_labels), labels.view(-1))
            elif self.config.problem_type == "multi_label_classification":
                loss = self.bce(logits, labels)

        result = DPLMSequenceClassifierOutput(
            loss=loss,
            logits=logits,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
            s_max=outputs.s_max,
        )
        return result if return_dict else result.to_tuple()


class DPLMForTokenClassification(DPLMPreTrainedModel, EmbeddingMixin):
    config_class = DPLMConfig

    def get_input_embeddings(self) -> nn.Module:
        return self.esm.get_input_embeddings()

    def set_input_embeddings(self, value: nn.Module) -> None:
        self.esm.set_input_embeddings(value)

    def __init__(self, config) -> None:
        DPLMPreTrainedModel.__init__(self, config)
        self.num_labels = config.num_labels
        self.esm = FAST_DPLM_ENCODER(config)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)
        self.classifier = nn.Linear(config.hidden_size, config.num_labels)
        self.loss_fct = nn.CrossEntropyLoss()
        self.post_init()

    def _embed(
        self,
        input_ids: torch.Tensor,
        attention_mask: torch.Tensor | None = None,
        hidden_state_index: int = -1,
        store_all_hidden_states: bool = False,
    ) -> torch.Tensor:
        return self.esm._embed(
            input_ids,
            attention_mask,
            hidden_state_index=hidden_state_index,
            store_all_hidden_states=store_all_hidden_states,
        )

    def forward(
        self,
        input_ids: torch.Tensor | None = None,
        attention_mask: torch.Tensor | None = None,
        inputs_embeds: torch.Tensor | None = None,
        labels: torch.Tensor | None = None,
        output_attentions: bool | None = None,
        output_hidden_states: bool | None = None,
        output_s_max: bool | None = False,
        return_dict: bool | None = None,
    ) -> tuple[torch.Tensor, ...] | DPLMTokenClassifierOutput:
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
        outputs = self.esm(
            input_ids=input_ids,
            attention_mask=attention_mask,
            inputs_embeds=inputs_embeds,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            output_s_max=output_s_max,
            return_dict=True,
        )
        sequence_output = self.dropout(outputs.last_hidden_state)  # (b, l, d)
        logits = self.classifier(sequence_output)  # (b, l, c)

        loss = None
        if labels is not None:
            labels = labels.to(logits.device)  # (b, l)
            loss = self.loss_fct(logits.view(-1, self.num_labels), labels.view(-1))  # ()

        result = DPLMTokenClassifierOutput(
            loss=loss,
            logits=logits,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
            s_max=outputs.s_max,
        )
        return result if return_dict else result.to_tuple()