FishCaduceus/FishCaduceus-28L-1024

1"""Reverse-complement equivariant modules.2 3"""4from collections import OrderedDict5from typing import Optional6 7import torch8from torch import Tensor9from torch import nn10from torch.nn import functional as F11 12try:13    from mamba_ssm.ops.triton.layernorm import RMSNorm, layer_norm_fn, rms_norm_fn  # Legacy mambav1 file structure14except ImportError:15    try:16        from mamba_ssm.ops.triton.layer_norm import RMSNorm, layer_norm_fn, rms_norm_fn  # mambav2 file structure17    except ImportError:18        RMSNorm, layer_norm_fn, rms_norm_fn = None, None, None19 20 21class RCPSEmbedding(nn.Module):22    """Embedding layer that supports reverse-complement equivariance."""23    def __init__(self, vocab_size: int, d_model: int, complement_map: dict, **factory_kwargs):24        """25        Args:26            vocab_size: Size of vocabulary.27            d_model: Dimensionality of embedding (actual embedding matrix will have 1/2 the output dim).28            complement_map: Dictionary mapping each token id to its complement.29        """30        super().__init__()31        self.register_buffer(32            "complement_map",33            torch.tensor(list(OrderedDict(complement_map).values()), dtype=torch.long)34        )35        self.embedding = nn.Embedding(vocab_size, d_model, **factory_kwargs)36 37    @property38    def weight(self):39        """Embedding weights."""40        return self.embedding.weight41 42    def set_weight(self, value):43        """Set embedding weights."""44        self.embedding.weight = value45 46    def rc(self, x):47        """Reverse-complement a tensor of input_ids by flipping along length dimension and complementing the ids."""48        return torch.gather(49            self.complement_map.unsqueeze(0).expand(x.shape[0], -1),50            dim=1,51            index=torch.flip(x, dims=[-1])52        )53 54    def forward(self, input_ids):55        """Reverse-complement equivariant forward pass.56 57        This embedding module doubles the output dimensionality to support reverse-complement equivariance.58 59        Args:60            input_ids: Input tensor of shape (batch_size, seq_len)61        Returns:62            Embedding tensor of shape (batch_size, seq_len, d_model * 2)63        """64        fwd_out = self.embedding(input_ids)65        rc_out = torch.flip(self.embedding(self.rc(input_ids)), dims=[-2, -1])66 67        return torch.cat([fwd_out, rc_out], dim=-1)68 69 70class RCPSWrapper(nn.Module):71    """Wrapper to convert arbitrary nn.Module into a reverse-complement equivariant module.72 73    See ref. "Towards a Better Understanding of Reverse-Complement Equivariance for Deep Learning Models in Regulatory74    Genomics", Zhou et al. (2022), https://proceedings.mlr.press/v165/zhou22a.html for more details.75    """76    def __init__(self, submodule: nn.Module):77        super().__init__()78        self.submodule = submodule79 80    @staticmethod81    def rc(x):82        """Reverse-complement a tensor by flipping the length (dim=-2) and channel (dim=-1) dimensions."""83        return torch.flip(x, dims=[-2, -1])84 85    def forward(self, x, **kwargs):86        """Reverse-complement equivariant forward pass.87 88        Args:89            x: Input tensor of shape (batch_size, seq_len, channels)90        Returns:91            Output tensor of shape (batch_size, seq_len, channels * 2)92        """93        n_channels = x.shape[-1]94        # Run submodule along sequence95        fwd_out = self.submodule(x[..., :n_channels // 2], **kwargs)96        # Run submodule along rc-sequence97        rc_out = self.submodule(self.rc(x[..., n_channels // 2:]), **kwargs)98        # Concatenate along channel dimension (dim=-1)99        return torch.cat([fwd_out, self.rc(rc_out)], dim=-1)100 101 102class RCPSAddNormWrapper(RCPSWrapper):103    """RC equivariant AddNorm layer."""104    def __init__(self, submodule: nn.Module):105        super().__init__(submodule)106 107    def forward(self, x, residual=None, prenorm=False):108        """109        Args:110            x: Input tensor of shape (batch_size, seq_len, channels)111            residual: Residual tensor of shape (batch_size, seq_len, channels) or None.112            prenorm: Whether to return residual.113        """114        n_channels = x.shape[-1]115        if residual is None:116            residual = x117            x_fwd = self.submodule(x[..., :n_channels // 2].to(dtype=self.submodule.weight.dtype))118            x_rc = self.submodule(self.rc(x[..., n_channels // 2:]).to(dtype=self.submodule.weight.dtype))119            x = torch.cat([x_fwd, self.rc(x_rc)], dim=-1)120        else:121            residual_fwd = x[..., :n_channels // 2] + residual[..., :n_channels // 2]122            x_fwd = self.submodule(residual_fwd.to(dtype=self.submodule.weight.dtype))123 124            residual_rc = self.rc(x[..., n_channels // 2:]) + self.rc(residual[..., n_channels // 2:])125            x_rc = self.submodule(residual_rc.to(dtype=self.submodule.weight.dtype))126 127            residual = torch.cat([residual_fwd, self.rc(residual_rc)], dim=-1)128            x = torch.cat([x_fwd, self.rc(x_rc)], dim=-1)129 130        return x if not prenorm else (x, residual)131 132 133class RCPSMambaBlock(nn.Module):134    def __init__(135            self,136            dim,137            mixer_cls,138            norm_cls=nn.LayerNorm,139            fused_add_norm=False,140            residual_in_fp32=False,141            device=None,  # Keep for consistency with original Mamba Block142            dtype=None,  # Keep for consistency with original Mamba Block143    ):144        """RCPS version of simple block wrapping a mixer class with LayerNorm/RMSNorm and residual connection.145 146        Adapted from: https://github.com/state-spaces/mamba/blob/main/mamba_ssm/modules/mamba_simple.py147        """148        super().__init__()149        self.residual_in_fp32 = residual_in_fp32150        self.fused_add_norm = fused_add_norm151        self.mixer = RCPSWrapper(mixer_cls(dim))152        norm_f = norm_cls(dim)153        self.norm = norm_f if fused_add_norm else RCPSAddNormWrapper(norm_f)154        if self.fused_add_norm:155            assert RMSNorm is not None, "RMSNorm import fails"156            assert isinstance(157                self.norm, (nn.LayerNorm, RMSNorm)158            ), "Only LayerNorm and RMSNorm are supported for fused_add_norm"159 160    def forward(161        self, hidden_states: Tensor, residual: Optional[Tensor] = None, inference_params=None162    ):163        r"""Pass the input through the encoder layer.164 165        Args:166            hidden_states: the sequence to the encoder layer (required).167            residual: hidden_states = Mixer(LN(residual)).168            inference_params: inference parameters for mixer.169        """170        if not self.fused_add_norm:171            hidden_states, residual = self.norm(hidden_states, residual=residual, prenorm=True)172            if self.residual_in_fp32:173                residual = residual.to(torch.float32)174        else:175            fused_add_norm_fn = rms_norm_fn if isinstance(self.norm, RMSNorm) else layer_norm_fn176 177            hidden_states_fwd, residual_fwd = fused_add_norm_fn(178                hidden_states[..., hidden_states.shape[-1] // 2:],179                self.norm.weight,180                self.norm.bias,181                residual=residual[..., hidden_states.shape[-1] // 2:] if residual is not None else None,182                prenorm=True,183                residual_in_fp32=self.residual_in_fp32,184                eps=self.norm.eps,185            )186 187            hidden_states_rc, residual_rc = fused_add_norm_fn(188                hidden_states[..., :hidden_states.shape[-1] // 2].flip(dims=[-2, -1]),189                self.norm.weight,190                self.norm.bias,191                residual=residual[..., :hidden_states.shape[-1] // 2].flip(dims=[-2, -1]) if residual is not None else None,192                prenorm=True,193                residual_in_fp32=self.residual_in_fp32,194                eps=self.norm.eps,195            )196            hidden_states = torch.cat([hidden_states_fwd, hidden_states_rc.flip(dims=[-2, -1])], dim=-1)197            residual = torch.cat([residual_fwd, residual_rc.flip(dims=[-2, -1])], dim=-1)198        hidden_states = self.mixer(hidden_states, inference_params=inference_params)199        return hidden_states, residual200 201    def allocate_inference_cache(self, batch_size, max_seqlen, dtype=None, **kwargs):202        """Allocate inference cache for mixer.203 204        Keep for compatibility with original Mamba Block.205        """206        return self.mixer.allocate_inference_cache(batch_size, max_seqlen, dtype=dtype, **kwargs)207 208 209class RCPSLMHead(nn.Module):210    """LM Head for reverse-complement equivariant inputs, which have dim * 2 relative to standard inputs."""211    def __init__(self, true_dim: int, vocab_size: int, complement_map: dict, **factory_kwargs):212        """213        `true_dim` corresponds to the actual dimensionality of the input were it not reverse-complement214        equivariant, i.e. 0.5 times the actual input dim.215        """216        super().__init__()217        self.register_buffer(218            "complement_map",219            torch.tensor(list(OrderedDict(complement_map).values()), dtype=torch.long)220        )221        self.true_dim = true_dim222        self.lm_head = nn.Linear(true_dim, vocab_size, bias=False, **factory_kwargs)223 224    @property225    def weight(self):226        """LM head weights."""227        return self.lm_head.weight228 229    def set_weight(self, value):230        """Set LM head weights."""231        self.lm_head.weight = value232 233    def forward(self, x):234        """235        Args:236            x: Input tensor of shape (batch_size, seq_len, dim), where dim = 2 * true_dim.237        """238        n_channels = x.shape[-1]239        assert n_channels == 2 * self.true_dim, "Input must have 2 * true_dim channels."240        fwd_logits = F.linear(x[..., :n_channels // 2], self.weight, bias=self.lm_head.bias)241        rc_logits = F.linear(242            torch.flip(x[..., n_channels // 2:], dims=[-1]),243            self.weight[self.complement_map, :],244            bias=self.lm_head.bias245        )246        return fwd_logits + rc_logits247