modeling_utrlm.py

"""UTR-LM ported to Hugging Face PreTrainedModel."""

import math
from typing import Optional, Tuple, Union

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import PreTrainedModel
from transformers.modeling_outputs import BaseModelOutput, MaskedLMOutput

from .configuration_utrlm import UtrLmConfig


# ---------------------------------------------------------------------------
# Rotary embeddings
# ---------------------------------------------------------------------------

def _rotate_half(x: torch.Tensor) -> torch.Tensor:
    x1, x2 = x.chunk(2, dim=-1)
    return torch.cat((-x2, x1), dim=-1)


def _apply_rotary_pos_emb(x, cos, sin):
    cos = cos[:, : x.shape[-2], :].to(x.dtype)
    sin = sin[:, : x.shape[-2], :].to(x.dtype)
    return (x * cos) + (_rotate_half(x) * sin)


class RotaryEmbedding(nn.Module):
    def __init__(self, dim: int):
        super().__init__()
        inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2).float() / dim))
        self.register_buffer("inv_freq", inv_freq)
        self._seq_len_cached: Optional[int] = None
        self._cos_cached: Optional[torch.Tensor] = None
        self._sin_cached: Optional[torch.Tensor] = None

    def _update_cos_sin_tables(self, x: torch.Tensor, seq_dimension: int = 1):
        seq_len = x.shape[seq_dimension]
        if seq_len != self._seq_len_cached or self._cos_cached.device != x.device:
            self._seq_len_cached = seq_len
            t = torch.arange(x.shape[seq_dimension], device=x.device).type_as(self.inv_freq)
            freqs = torch.einsum("i,j->ij", t, self.inv_freq)
            emb = torch.cat((freqs, freqs), dim=-1).to(x.device)
            self._cos_cached = emb.cos()[None, :, :]
            self._sin_cached = emb.sin()[None, :, :]
        return self._cos_cached, self._sin_cached

    def forward(self, q, k):
        self._cos_cached, self._sin_cached = self._update_cos_sin_tables(k, seq_dimension=-2)
        return (
            _apply_rotary_pos_emb(q, self._cos_cached, self._sin_cached),
            _apply_rotary_pos_emb(k, self._cos_cached, self._sin_cached),
        )


# ---------------------------------------------------------------------------
# Attention variants
# ---------------------------------------------------------------------------

class UtrLmAttention(nn.Module):
    """Eager (standard) attention."""

    def __init__(self, embed_dim: int, num_heads: int):
        super().__init__()
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads
        self.scaling = self.head_dim ** -0.5

        self.k_proj = nn.Linear(embed_dim, embed_dim)
        self.v_proj = nn.Linear(embed_dim, embed_dim)
        self.q_proj = nn.Linear(embed_dim, embed_dim)
        self.out_proj = nn.Linear(embed_dim, embed_dim)
        self.rot_emb = RotaryEmbedding(dim=self.head_dim)

    def _project(self, x):
        """Project and reshape x (T, B, E) -> q/k/v in (B*H, T, head_dim)."""
        tgt_len, bsz, _ = x.size()
        q = (self.q_proj(x) * self.scaling).contiguous().view(tgt_len, bsz * self.num_heads, self.head_dim).transpose(0, 1)
        k = self.k_proj(x).contiguous().view(tgt_len, bsz * self.num_heads, self.head_dim).transpose(0, 1)
        v = self.v_proj(x).contiguous().view(tgt_len, bsz * self.num_heads, self.head_dim).transpose(0, 1)
        q, k = self.rot_emb(q, k)
        return q, k, v

    def forward(self, x, key_padding_mask, output_attentions: bool = False):
        tgt_len, bsz, _ = x.size()
        q, k, v = self._project(x)

        attn_weights = torch.bmm(q, k.transpose(1, 2))
        attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, tgt_len)
        if key_padding_mask is not None:
            attn_weights = attn_weights.masked_fill(
                key_padding_mask.unsqueeze(1).unsqueeze(2).to(torch.bool), float("-inf")
            )
        attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, tgt_len)

        attn_probs = F.softmax(attn_weights, dim=-1, dtype=torch.float32).type_as(attn_weights)
        attn = torch.bmm(attn_probs, v)
        attn = attn.transpose(0, 1).contiguous().view(tgt_len, bsz, self.embed_dim)
        out = self.out_proj(attn)

        if output_attentions:
            return out, attn_probs.view(bsz, self.num_heads, tgt_len, tgt_len)
        return out, None


class UtrLmSdpaAttention(UtrLmAttention):
    """SDPA attention via torch.nn.functional.scaled_dot_product_attention."""

    def forward(self, x, key_padding_mask, output_attentions: bool = False):
        if output_attentions:
            # SDPA doesn't expose attention weights; fall back to eager.
            return super().forward(x, key_padding_mask, output_attentions=True)

        tgt_len, bsz, _ = x.size()
        q, k, v = self._project(x)  # (B*H, T, head_dim)

        # Reshape to (B, H, T, head_dim) for SDPA
        q = q.view(bsz, self.num_heads, tgt_len, self.head_dim)
        k = k.view(bsz, self.num_heads, tgt_len, self.head_dim)
        v = v.view(bsz, self.num_heads, tgt_len, self.head_dim)

        # Convert bool padding mask -> additive float mask (B, 1, 1, T)
        attn_mask = None
        if key_padding_mask is not None:
            attn_mask = torch.zeros(bsz, 1, 1, tgt_len, dtype=q.dtype, device=q.device)
            attn_mask = attn_mask.masked_fill(key_padding_mask[:, None, None, :], float("-inf"))

        # scale=1.0 because q is already pre-scaled by self.scaling
        out = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask, scale=1.0)
        out = out.permute(2, 0, 1, 3).contiguous().view(tgt_len, bsz, self.embed_dim)
        return self.out_proj(out), None


class UtrLmFlashAttention2(UtrLmAttention):
    """Flash Attention 2 via flash_attn (must be installed separately)."""

    def forward(self, x, key_padding_mask, output_attentions: bool = False):
        if output_attentions:
            # Flash attention doesn't expose attention weights; fall back to eager.
            return super().forward(x, key_padding_mask, output_attentions=True)

        try:
            from flash_attn import flash_attn_func
            from flash_attn.bert_padding import pad_input, unpad_input
        except ImportError as e:
            raise ImportError("flash_attn is required for attn_implementation='flash_attention_2'. "
                              "Install with: pip install flash-attn --no-build-isolation") from e

        tgt_len, bsz, _ = x.size()
        q, k, v = self._project(x)  # (B*H, T, head_dim)

        # Reshape to (B, T, H, head_dim) - flash_attn's expected layout
        q = q.view(bsz, self.num_heads, tgt_len, self.head_dim).permute(0, 2, 1, 3)
        k = k.view(bsz, self.num_heads, tgt_len, self.head_dim).permute(0, 2, 1, 3)
        v = v.view(bsz, self.num_heads, tgt_len, self.head_dim).permute(0, 2, 1, 3)

        # Flash attention requires fp16 or bf16
        orig_dtype = q.dtype
        if orig_dtype not in (torch.float16, torch.bfloat16):
            q, k, v = q.to(torch.bfloat16), k.to(torch.bfloat16), v.to(torch.bfloat16)

        if key_padding_mask is not None:
            # Unpad, run varlen flash attention, repad
            from flash_attn import flash_attn_varlen_func
            attention_mask = ~key_padding_mask  # True = valid token
            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)

            out_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,
                softmax_scale=1.0,  # q already pre-scaled
                causal=False,
            )
            out = pad_input(out_unpad, indices, bsz, tgt_len)
        else:
            out = flash_attn_func(q, k, v, softmax_scale=1.0, causal=False)

        out = out.to(orig_dtype).permute(1, 0, 2, 3).contiguous().view(tgt_len, bsz, self.embed_dim)
        return self.out_proj(out), None


UTRLM_ATTENTION_CLASSES = {
    "eager": UtrLmAttention,
    "sdpa": UtrLmSdpaAttention,
    "flash_attention_2": UtrLmFlashAttention2,
}


# ---------------------------------------------------------------------------
# Transformer layer (pre-LN)
# ---------------------------------------------------------------------------

def _gelu(x):
    return x * 0.5 * (1.0 + torch.erf(x / math.sqrt(2.0)))


class UtrLmLayer(nn.Module):
    def __init__(self, embed_dim: int, attention_heads: int, config: UtrLmConfig):
        super().__init__()
        attn_cls = UTRLM_ATTENTION_CLASSES[getattr(config, "_attn_implementation", "eager")]
        self.self_attn = attn_cls(embed_dim, attention_heads)
        self.self_attn_layer_norm = nn.LayerNorm(embed_dim)
        self.fc1 = nn.Linear(embed_dim, 4 * embed_dim)
        self.fc2 = nn.Linear(4 * embed_dim, embed_dim)
        self.final_layer_norm = nn.LayerNorm(embed_dim)

    def forward(self, x, padding_mask, output_attentions: bool = False):
        residual = x
        x = self.self_attn_layer_norm(x)
        x, attn_weights = self.self_attn(x, key_padding_mask=padding_mask, output_attentions=output_attentions)
        x = residual + x

        residual = x
        x = self.final_layer_norm(x)
        x = _gelu(self.fc1(x))
        x = self.fc2(x)
        return residual + x, attn_weights


# ---------------------------------------------------------------------------
# Backbone
# ---------------------------------------------------------------------------

class UtrLmModel(PreTrainedModel):
    """
    UTR-LM encoder backbone. Returns last_hidden_state (B, T, E).
    The [CLS] token sits at position 0 (prepend_bos=True by default).
    """

    config_class = UtrLmConfig
    base_model_prefix = "utrlm"
    _supports_sdpa = True
    _supports_flash_attn_2 = True

    def __init__(self, config: UtrLmConfig):
        super().__init__(config)
        self.embed_scale = 1
        self.embed_tokens = nn.Embedding(
            config.alphabet_size, config.embed_dim, padding_idx=config.padding_idx
        )
        self.layers = nn.ModuleList(
            [UtrLmLayer(config.embed_dim, config.attention_heads, config) for _ in range(config.num_layers)]
        )
        self.emb_layer_norm_after = nn.LayerNorm(config.embed_dim)
        self.post_init()

    def get_input_embeddings(self):
        return self.embed_tokens

    def set_input_embeddings(self, value):
        self.embed_tokens = value

    def forward(
        self,
        input_ids: torch.LongTensor,
        attention_mask: Optional[torch.BoolTensor] = None,
        output_hidden_states: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[Tuple, BaseModelOutput]:
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None
            else self.config.output_hidden_states
        )
        output_attentions = (
            output_attentions if output_attentions is not None else self.config.output_attentions
        )
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        cfg = self.config
        # HF convention: attention_mask is 1=attend, 0=pad.
        # Convert to bool padding_mask (True = ignore) or derive from input_ids.
        if attention_mask is not None:
            padding_mask = attention_mask.eq(0)
        else:
            padding_mask = input_ids.eq(cfg.padding_idx)

        x = self.embed_scale * self.embed_tokens(input_ids)

        if cfg.token_dropout:
            x.masked_fill_((input_ids == cfg.mask_idx).unsqueeze(-1), 0.0)
            mask_ratio_train = 0.15 * 0.8
            src_lengths = (~padding_mask).sum(-1)
            mask_ratio_observed = (input_ids == cfg.mask_idx).sum(-1).to(x.dtype) / src_lengths.to(x.dtype)
            x = x * (1 - mask_ratio_train) / (1 - mask_ratio_observed)[:, None, None]

        if padding_mask is not None:
            x = x * (1 - padding_mask.unsqueeze(-1).type_as(x))

        all_hidden_states = () if output_hidden_states else None
        all_attentions = () if output_attentions else None
        if output_hidden_states:
            all_hidden_states += (x,)

        x = x.transpose(0, 1)  # (B, T, E) -> (T, B, E)
        effective_padding = padding_mask if padding_mask.any() else None

        for layer in self.layers:
            x, attn_weights = layer(x, padding_mask=effective_padding, output_attentions=output_attentions)
            if output_hidden_states:
                all_hidden_states += (x.transpose(0, 1),)
            if output_attentions:
                all_attentions += (attn_weights,)

        x = self.emb_layer_norm_after(x)
        x = x.transpose(0, 1)  # (T, B, E) -> (B, T, E)

        if output_hidden_states:
            all_hidden_states = all_hidden_states[:-1] + (x,)

        if not return_dict:
            return tuple(v for v in [x, all_hidden_states, all_attentions] if v is not None)

        return BaseModelOutput(
            last_hidden_state=x,
            hidden_states=all_hidden_states,
            attentions=all_attentions,
        )


# ---------------------------------------------------------------------------
# MLM head
# ---------------------------------------------------------------------------

class UtrLmForMaskedLM(PreTrainedModel):
    """
    UTR-LM with a masked-language-modelling head.
    Returns MaskedLMOutput with logits (B, T, vocab_size).
    """

    config_class = UtrLmConfig
    base_model_prefix = "utrlm"
    _supports_sdpa = True
    _supports_flash_attn_2 = True

    def __init__(self, config: UtrLmConfig):
        super().__init__(config)
        self.utrlm = UtrLmModel(config)

        embed_dim = config.embed_dim
        vocab_size = config.alphabet_size
        self.lm_head = nn.ModuleDict({
            "dense": nn.Linear(embed_dim, embed_dim),
            "layer_norm": nn.LayerNorm(embed_dim),
        })
        self.lm_head_bias = nn.Parameter(torch.zeros(vocab_size))

        self.post_init()

    def get_input_embeddings(self):
        return self.utrlm.embed_tokens

    def set_input_embeddings(self, value):
        self.utrlm.embed_tokens = value

    def get_output_embeddings(self):
        return self.utrlm.embed_tokens

    def set_output_embeddings(self, new_embeddings):
        self.utrlm.embed_tokens = new_embeddings

    def _lm_head_forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.lm_head["dense"](x)
        x = _gelu(x)
        x = self.lm_head["layer_norm"](x)
        return F.linear(x, self.utrlm.embed_tokens.weight) + self.lm_head_bias

    def forward(
        self,
        input_ids: torch.LongTensor,
        attention_mask: Optional[torch.BoolTensor] = None,
        labels: Optional[torch.LongTensor] = None,
        output_hidden_states: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[Tuple, MaskedLMOutput]:
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        outputs = self.utrlm(
            input_ids,
            attention_mask=attention_mask,
            output_hidden_states=output_hidden_states,
            output_attentions=output_attentions,
            return_dict=True,
        )
        logits = self._lm_head_forward(outputs.last_hidden_state)

        loss = None
        if labels is not None:
            loss = F.cross_entropy(
                logits.view(-1, self.config.alphabet_size),
                labels.view(-1),
                ignore_index=self.config.padding_idx,
            )

        if not return_dict:
            output = (logits,) + outputs[1:]
            return (loss,) + output if loss is not None else output

        return MaskedLMOutput(
            loss=loss,
            logits=logits,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
        )