Delete sct.py

sct.py

@@ -1,756 +0,0 @@
-import math
-from dataclasses import dataclass
-from typing import Optional, Tuple
-
-import numpy as np
-import torch
-import torch.nn as nn
-import torch.nn.functional as F  # noqa: N812
-from transformers import PretrainedConfig, PreTrainedModel
-
-
-class GeLU(nn.Module):
-    def __init__(self) -> None:
-        """
-        This is the gelu implementation from the original ESM repo.
-        Using F.gelu yields subtly wrong results.
-        """
-        super().__init__()
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        return x * 0.5 * (1.0 + torch.erf(x / math.sqrt(2.0)))
-
-
-@dataclass
-class RotaryEmbeddingConfig:
-    """
-    Parameters to initialize the RotaryEmbedding layer. The rescaling factor allows
-    to adapt the rotary embeddings to larger lengths than what was used for training.
-    One of this strategy is presented in the Yarn paper: https://arxiv.org/pdf/2309.00071.pdf. # noqa
-    Args:
-    """
-
-    rescaling_factor: Optional[float]
-
-
-class RotaryEmbedding(torch.nn.Module):
-    """
-    Rotary position embeddings based on those in
-    [RoFormer](https://huggingface.co/docs/transformers/model_doc/roformer).
-    Query and keys are transformed by rotation
-    matrices which depend on their relative positions.
-    """
-
-    def __init__(self, dim: int, rotary_embedding_config: RotaryEmbeddingConfig):
-        super().__init__()
-
-        # Extract argument from the config
-        self.rescaling_factor = rotary_embedding_config.rescaling_factor
-        self.upper_freq = 10000
-        self.dim = dim
-
-        self._seq_len_cached = None
-        self._cos_cached = None
-        self._sin_cached = None
-
-    def _apply_rotary_pos_emb(
-        self,
-        heads: torch.Tensor,
-        cos: torch.Tensor,
-        sin: torch.Tensor,
-    ) -> torch.Tensor:
-        """ """
-        x_first, x_second = (
-            heads[..., : heads.shape[-1] // 2],
-            heads[..., heads.shape[-1] // 2 :],
-        )
-
-        first_part = x_first * cos - x_second * sin
-        second_part = x_second * cos + x_first * sin
-
-        return torch.cat((first_part, second_part), dim=-1)
-
-    def _compute_cos_sin_tables(
-        self, x: torch.Tensor, inv_freq: torch.Tensor, seq_dimension: int = 2
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        seq_len = x.shape[seq_dimension]
-        # Reset the tables if the sequence length has changed,
-        # or if we're on a new device (possibly due to tracing for instance)
-        self._seq_len_cached = seq_len
-        t = torch.arange(x.shape[seq_dimension], device=x.device).type_as(inv_freq)
-        # freqs = torch.outer(t, inv_freq)
-        freqs = torch.einsum("i, j -> ij", t, inv_freq)
-
-        self._cos_cached = torch.cos(freqs)[None, :, None, :]
-        self._sin_cached = torch.sin(freqs)[None, :, None, :]
-        # emb = torch.cat((freqs, freqs), dim=-1).to(x.device)
-
-        # self._cos_cached = emb.cos()[None, None, :, :]
-        # self._sin_cached = emb.sin()[None, None, :, :]
-
-        return self._cos_cached, self._sin_cached
-
-    def forward(
-        self, q: torch.Tensor, k: torch.Tensor
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
-        if self.rescaling_factor is None:
-            inv_freq = 1.0 / (
-                self.upper_freq ** (torch.arange(0, self.dim, 2).float() / self.dim)
-            )
-        else:
-            updated_base = self.upper_freq * (
-                self.rescaling_factor ** (self.dim / (self.dim - 2))
-            )
-            inv_freq = 1.0 / (
-                updated_base ** (torch.arange(0, self.dim, 2).float() / self.dim)
-            )
-
-        self._cos_cached, self._sin_cached = self._compute_cos_sin_tables(
-            q,
-            inv_freq,
-            seq_dimension=-3,
-        )
-
-        return (
-            self._apply_rotary_pos_emb(q, self._cos_cached, self._sin_cached),
-            self._apply_rotary_pos_emb(k, self._cos_cached, self._sin_cached),
-        )
-
-
-class ResidualConvBlock(nn.Module):
-    """
-    Conv Block with Residual connection.
-    """
-
-    def __init__(self, dim_in: int, dim_out: int, seq_len: int, kernel_size: int = 1):
-        super().__init__()
-        self.conv_block = ConvBlock(
-            dim_in=dim_in, dim_out=dim_out, seq_len=seq_len, kernel_size=kernel_size
-        )
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        y = self.conv_block(x)
-        return x.reshape(y.shape) + y
-
-
-class ConvBlock(nn.Module):
-    """
-    Conv Block.
-    """
-
-    def __init__(self, dim_in: int, dim_out: int, seq_len: int, kernel_size: int = 1):
-        super().__init__()
-        self.conv = nn.Conv1d(
-            in_channels=dim_in,
-            out_channels=dim_out,
-            kernel_size=kernel_size,
-            padding="same",
-        )
-        self.layer_norm = nn.LayerNorm(seq_len, eps=1e-5)
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        x = self.layer_norm(x)
-        x = x.reshape(x.shape[0], x.shape[1], -1)
-        x = self.conv(x)
-        x = F.gelu(x, approximate="tanh")
-        return x
-
-
-class ResidualDeConvBlock(nn.Module):
-    """
-    Conv Block with Residual connection.
-    """
-
-    def __init__(
-        self,
-        dim_in: int,
-        dim_out: int,
-        seq_len: int,
-        kernel_size: int = 1,
-        stride: int = 1,
-    ):
-        super().__init__()
-        self.deconv_block = DeConvBlock(
-            dim_in=dim_in,
-            dim_out=dim_out,
-            seq_len=seq_len,
-            kernel_size=kernel_size,
-            stride=stride,
-        )
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        y = self.deconv_block(x)
-        return x.reshape(y.shape) + y
-
-
-class DeConvBlock(nn.Module):
-    """
-    DeConv Block.
-    """
-
-    def __init__(
-        self,
-        dim_in: int,
-        dim_out: int,
-        seq_len: int,
-        kernel_size: int = 1,
-        stride: int = 1,
-    ):
-        super().__init__()
-        self.deconv = nn.ConvTranspose1d(
-            in_channels=dim_in,
-            out_channels=dim_out,
-            kernel_size=kernel_size,
-            stride=stride,
-            padding=0,
-        )
-        self.layer_norm = nn.LayerNorm(seq_len)
-        self.kernel_size = kernel_size
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        x = self.layer_norm(x)
-        x = x.reshape(x.shape[0], x.shape[1], -1)
-        x = self.deconv(x)
-        if self.kernel_size == 5:
-            # handle the special case where haiku
-            # deconv removes padding automatically
-            x = x[:, :, 1:-2]
-        x = F.gelu(x, approximate="tanh")
-        return x
-
-
-class SpatialEncoding(nn.Module):
-    """
-    Spatial coordinates encoding module
-    """
-
-    def __init__(
-        self,
-        embed_dim: int,
-        num_scales: int = 10,
-        sigma_min: float = 1.0,
-        sigma_max: float = 10.0,
-    ):
-        super().__init__()
-        self.num_scales = num_scales
-        self.sigma_min = sigma_min
-        self.sigma_max = sigma_max
-        self.g = sigma_max / sigma_min
-        self.scales = torch.linspace(sigma_min, sigma_max, num_scales)
-        self.fc_layer = nn.Linear(embed_dim, embed_dim)
-
-    def scale_specific_encoder(
-        self, coordinates: torch.Tensor, scale: float
-    ) -> torch.Tensor:
-        x, y = coordinates[..., 0], coordinates[..., 1]
-        constant = self.sigma_min * (self.g ** (scale / (self.num_scales - 1)))
-        x_transform = torch.cos(x / constant)
-        y_transform = torch.sin(y / constant)
-        transformed_coordinates = torch.stack([x_transform, y_transform], dim=-1)
-        return transformed_coordinates
-
-    def forward(self, coordinates: torch.Tensor) -> torch.Tensor:
-        transformed_coordinates = [
-            self.scale_specific_encoder(coordinates, scale) for scale in self.scales
-        ]
-        transformed_coordinates = torch.cat(transformed_coordinates, dim=-1)
-        return self.fc_layer(transformed_coordinates)
-
-
-class ConvTowerBlock(nn.Module):
-    def __init__(
-        self, dim_in: int, dim_out: int, seq_len: int, kernel_size: int, num_cells: int
-    ) -> None:
-        super().__init__()
-        self.conv_layer = ConvBlock(
-            dim_in=dim_in, dim_out=dim_out, seq_len=seq_len, kernel_size=kernel_size
-        )
-        self.res_conv = ResidualConvBlock(
-            dim_in=dim_out, dim_out=dim_out, seq_len=seq_len, kernel_size=1
-        )
-        self.avg_pool = nn.AvgPool1d(kernel_size=2, stride=2)
-        self.num_cells = num_cells
-
-    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
-        residual = x
-        x = x.reshape(x.shape[0], x.shape[1], self.num_cells, -1)  # noqa: FKA100
-        x = self.conv_layer(x)
-        x = x.reshape((x.shape[0], x.shape[1], self.num_cells, -1))
-        x = self.res_conv(x)
-        x = self.avg_pool(x)
-        return x, residual
-
-
-class DeConvTowerBlock(nn.Module):
-    def __init__(
-        self,
-        dim_in: int,
-        dim_out: int,
-        kernel_size: int,
-        seq_len: int,
-        stride: int = 2,
-        num_cells: int = 1,
-    ):
-        super().__init__()
-        self.deconv_block = DeConvBlock(
-            dim_in=dim_in,
-            dim_out=dim_out,
-            seq_len=seq_len,
-            kernel_size=kernel_size,
-            stride=stride,
-        )
-        self.res_deconv_block = ResidualDeConvBlock(
-            dim_in=dim_out, dim_out=dim_out, seq_len=seq_len * 2, kernel_size=1
-        )
-        self.num_cells = num_cells
-
-    def forward(self, x: torch.Tensor, res: torch.Tensor) -> torch.Tensor:
-        x = x.reshape((x.shape[0], x.shape[1], self.num_cells, -1))
-        x = self.deconv_block(x)
-        x = x.reshape((x.shape[0], x.shape[1], self.num_cells, -1))
-        x = self.res_deconv_block(x)
-
-        x = x + res
-        return x
-
-
-class MultiHeadAttention(nn.Module):
-    def __init__(
-        self,
-        num_heads: int,
-        key_size: int,
-        rotary_embedding_config: Optional[RotaryEmbeddingConfig] = None,
-        add_bias_kv: bool = False,
-        value_size: Optional[int] = None,
-        model_size: Optional[int] = None,
-        name: Optional[str] = None,
-    ):
-        super().__init__()
-        if not model_size:
-            model_size = key_size
-        if not value_size:
-            value_size = key_size
-        self.model_size = model_size
-        self.key_size = key_size
-        self.value_size = value_size
-        self.add_bias_kv = add_bias_kv
-        self.name = name
-        self.num_heads = num_heads
-        self._rotary_embedding_config = rotary_embedding_config
-
-        self.w_k = nn.Linear(self.model_size, self.num_heads * self.key_size)
-        self.w_q = nn.Linear(self.model_size, self.num_heads * self.key_size)
-        self.w_v = nn.Linear(self.model_size, self.num_heads * self.value_size)
-        self.output = nn.Linear(self.num_heads * self.value_size, self.model_size)
-        if self._rotary_embedding_config:
-            self._rotary_embedding = RotaryEmbedding(
-                self.key_size, self._rotary_embedding_config
-            )
-
-    def apply_rotary_embeddings(
-        self,
-        query: torch.Tensor,
-        key: torch.Tensor,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        """ """
-        query, key = self._rotary_embedding(query, key)
-        return query, key
-
-    def forward(
-        self,
-        query: torch.Tensor,
-        key: torch.Tensor,
-        value: torch.Tensor,
-        attention_mask: torch.Tensor | None = None,
-        attention_weight_bias: torch.Tensor | None = None,
-    ) -> dict[str, torch.Tensor]:
-        """
-        Returns:
-            dictionary containing attention weights
-            and outputs.
-        """
-        key_heads = self.w_k(key).reshape(
-            (*key.shape[:-1], self.num_heads, self.key_size)
-        )
-        query_heads = self.w_q(query).reshape(
-            (*query.shape[:-1], self.num_heads, self.key_size)
-        )
-        value_heads = self.w_v(value).reshape(
-            (*value.shape[:-1], self.num_heads, self.value_size)
-        )
-        if self._rotary_embedding_config:
-            query_heads, key_heads = self.apply_rotary_embeddings(
-                query_heads, key_heads
-            )
-        attention_weights = torch.einsum(
-            "...thd, ...Thd -> ...htT", query_heads, key_heads
-        )
-        sqrt_key_size = np.sqrt(self.key_size)
-        attention_weights = attention_weights / sqrt_key_size
-        if attention_mask:
-            attention_weights = torch.where(attention_mask, attention_weights, -1e30)
-        if attention_weight_bias:
-            attention_weights = F.softmax(
-                attention_weights + attention_weight_bias, dim=-1
-            )
-        else:
-            attention_weights = F.softmax(attention_weights, dim=-1)
-        value_out = torch.einsum(
-            "...htT, ...Thd->...thd", attention_weights, value_heads
-        )
-        value_out = value_out.reshape((*value_out.shape[:-2], -1))
-        embeddings = self.output(value_out)
-
-        return {"attention_weights": attention_weights, "embeddings": embeddings}
-
-
-class SelfAttentionBlock(nn.Module):
-    def __init__(
-        self,
-        num_heads: int,
-        embed_dim: int,
-        ffn_embed_dim: int,
-        key_size: Optional[int] = None,
-        add_bias_kv: bool = False,
-        add_bias_fnn: bool = True,
-        ffn_activation_name: str = "gelu-no-approx",
-        use_glu_in_ffn: bool = False,
-        layer_norm_eps: float = 1e-5,  # this is the default haiku value
-        pre_layer_norm: bool = True,
-        name: Optional[str] = None,
-        rotary_embedding_config: Optional[RotaryEmbeddingConfig] = None,
-    ):
-        super().__init__()
-        if key_size is None:
-            if embed_dim % num_heads != 0:
-                raise ValueError(
-                    f"The embedding dimension should be divisible by the number of "
-                    f"heads, however provided embedding dimension is {embed_dim} and "
-                    f"the number of heads is {num_heads}."
-                )
-            else:
-                key_size = embed_dim // num_heads
-
-        # Get ffn activation function
-        self._pre_layer_norm = pre_layer_norm
-        self._use_glu_in_fnn = use_glu_in_ffn
-        # Define layers
-        if use_glu_in_ffn:
-            # user should multiply ffn_embed_dim by 2/3 when using GLU
-            # to keep total number of parameters equal
-            # see https://arxiv.org/pdf/2002.05202.pdf. for more details
-            # we multiply by 2 here as the output will be split in 2 for GLU
-            self.fc1 = nn.Linear(embed_dim, int(2 * ffn_embed_dim), bias=add_bias_fnn)
-        else:
-            self.fc1 = nn.Linear(embed_dim, ffn_embed_dim, bias=add_bias_fnn)
-
-        self.fc2 = nn.Linear(ffn_embed_dim, embed_dim, bias=add_bias_fnn)
-
-        self.layer_norm_self_attention = nn.LayerNorm(
-            embed_dim,
-        )
-        self.layer_norm_mlp = nn.LayerNorm(embed_dim)
-        if ffn_activation_name == "swish":
-            self._ffn_activation_fn = nn.SiLU()
-        elif ffn_activation_name == "gelu-no-approx":
-            self._ffn_activation_fn = nn.GeLU(approximate="tanh")
-        else:
-            self._ffn_activation_fn = getattr(torch.nn, ffn_activation_name)
-
-        self.mha = MultiHeadAttention(
-            num_heads=num_heads,
-            key_size=key_size,
-            add_bias_kv=add_bias_kv,
-            model_size=embed_dim,
-            name="self_attention",
-            rotary_embedding_config=rotary_embedding_config,
-        )
-
-    def mlp(self, embed: torch.Tensor) -> torch.Tensor:
-
-        if self._pre_layer_norm:
-            x = self.layer_norm_mlp(embed)
-        else:
-            x = embed
-
-        if self._use_glu_in_fnn:
-            x = self.fc1(x)
-            x1, x2 = torch.split(x, split_size_or_sections=x.shape[-1] // 2, dim=-1)
-            x = self._ffn_activation_fn(x1) * x2
-        else:
-            x = self._ffn_activation_fn(self.fc1(x))
-        x = self.fc2(x)
-
-        if not self._pre_layer_norm:
-            x = self.layer_norm_mlp(x + embed)
-        return x
-
-    def forward(
-        self,
-        x: torch.Tensor,
-        attention_mask: torch.Tensor | None = None,
-        attention_weight_bias: torch.Tensor | None = None,
-    ) -> torch.Tensor:
-
-        res = x
-        if self._pre_layer_norm:
-            x = self.layer_norm_self_attention(x)
-
-        output = self.mha(
-            x,
-            x,
-            x,
-            attention_mask=attention_mask,
-            attention_weight_bias=attention_weight_bias,
-        )
-
-        if not self._pre_layer_norm:
-            output["embeddings"] = self.layer_norm_self_attention(
-                output["embeddings"] + res
-            )
-
-            x = output["embeddings"]
-        else:
-            x = output["embeddings"]
-            x = res + x
-
-        # MLP
-        if not self._pre_layer_norm:
-            x = self.mlp(x)
-        else:
-            x = x + self.mlp(x)
-
-        output["embeddings"] = x
-        return output
-
-
-class LMHead(nn.Module):
-    def __init__(
-        self, dim_in: int, embed_dim: int, dim_out: int, num_hidden_layers: int
-    ) -> None:
-        """ """
-        super().__init__()
-        self.num_hidden_layers = num_hidden_layers
-        self.linear_layers = nn.ModuleList([nn.Linear(dim_in, embed_dim)])
-        self.linear_layers.extend(
-            nn.ModuleList(
-                [nn.Linear(embed_dim, embed_dim)] for _ in range(num_hidden_layers - 1)
-            )
-        )
-        self.linear_out = nn.Linear(embed_dim, dim_out)
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        res = x  # noqa: F841
-        x = F.gelu(x, approximate="tanh")
-        for layer in self.linear_layers:
-            x = layer(x)
-            x = F.gelu(x, approximate="tanh")
-        out = self.linear_out(x)
-        return out
-
-
-@dataclass
-class sCTConfig(PretrainedConfig):  # noqa: N801
-    model_type = "sCT"
-
-    def __init__(self, **kwargs):  # type: ignore
-        self.alphabet_size = kwargs.get("alphabet_size", 7)
-        self.pad_token_id = kwargs.get("pad_token_id", 5)
-        self.mask_token_id = kwargs.get("mask_token_id", 6)
-        self.cell_len = kwargs.get("cell_len", 19968)
-
-        self.num_downsamples = kwargs.get("num_downsamples", 8)
-        self.attention_heads = kwargs.get("attention_heads", 16)
-        self.key_size = kwargs.get("key_size", None)
-        self.token_embed_dim = kwargs.get("token_embed_dim", 16)
-
-        self.embed_dim = kwargs.get("embed_dim", 1024)
-        self.ffn_embed_dim = kwargs.get("ffn_embed_dim", 2048)
-        self.num_layers = kwargs.get("num_layers", 4)
-        self.layer_norm_eps = kwargs.get("layer_norm_eps", 1e-5)
-        self.interpolation_method = kwargs.get("interpolation_method", "nearest")
-
-        # bad hack to satisfy cellnt_celltype_annotation.py:312
-        self.max_positions: int = kwargs.get("max_positions", 20480)
-        self.num_cells: int = kwargs.get("num_cells", 50)
-        self.num_hidden_layers_head: int = kwargs.get("num_hidden_layers_head", 1)
-
-        self.use_skip_connection: bool = kwargs.get("use_skip_connection", True)
-
-        # logging
-        self.use_gradient_checkpointing: bool = False
-
-        # return
-        self.embeddings_layers_to_save: Tuple[int, ...] = kwargs.get(
-            "embeddings_layers_to_save", ()
-        )
-        self.attention_maps_to_save: list[tuple[int, int]] = kwargs.get(
-            "attention_maps_to_save", []
-        )
-
-        # Spatial info configuration
-        self.use_spatial_information: bool = kwargs.get(
-            "use_spatial_information", False
-        )
-        self.num_scales: int = kwargs.get("num_scales", 10)
-        self.sigma_min: float = kwargs.get("sigma_min", 1.0)
-        self.sigma_max: float = kwargs.get("sigma_max", 10.0)
-
-        super().__init__(**kwargs)
-
-        def __post_init__(self) -> None:  # type: ignore # noqa: N807
-            """
-            Checks that the given values are compatible.
-            """
-            if self.key_size is None:
-                if not self.embed_dim % self.attention_heads == 0:
-                    raise ValueError(
-                        f"When no key size is provided, the embedding dimension"
-                        f"should be divisible by the number of heads, however "
-                        f"provided embedding dimension is {self.embed_dim} and "
-                        f"the number of heads is {self.attention_heads}."
-                    )
-                self.key_size = self.embed_dim // self.attention_heads
-
-
-class sCT(PreTrainedModel):  # noqa: N801
-    config_class = sCTConfig
-
-    def __init__(self, config: sCTConfig):
-        # super().__init__(config)
-        super().__init__(config=config)
-        if config.use_spatial_information:
-            self.spatial_embed_layer = SpatialEncoding(
-                embed_dim=config.token_embed_dim,
-                num_scales=config.num_scales,
-                sigma_min=config.sigma_min,
-                sigma_max=config.sigma_max,
-            )
-        self.cell_len = config.cell_len
-
-        self.token_embed = nn.Embedding(config.alphabet_size, config.token_embed_dim)
-
-        attention_maps_to_save = config.attention_maps_to_save
-        self._attention_layers_to_save = list({t[0] for t in attention_maps_to_save})
-
-        self._attention_maps_per_layer_to_save = {
-            layer: [t[1] for t in attention_maps_to_save if t[0] == layer]
-            for layer in self._attention_layers_to_save
-        }
-
-        max_layer = max(self._attention_layers_to_save + [0])
-        if max_layer > config.num_layers:
-            raise ValueError(
-                f"You are requiring attention maps for layer {max_layer}, "
-                f"while the model has {config.num_layers} layers only."
-            )
-
-        filter_list = np.linspace(
-            config.token_embed_dim,
-            config.embed_dim,
-            config.num_downsamples + 1,
-        )
-
-        filter_list = np.ceil(filter_list / 32) * 32
-        filter_list = filter_list.astype(int).tolist()
-
-        self._filter_list = filter_list
-        self._rotary_embedding_config = RotaryEmbeddingConfig(rescaling_factor=None)
-
-        self.stem_conv = nn.Sequential(
-            nn.Conv1d(
-                in_channels=config.token_embed_dim,
-                out_channels=config.token_embed_dim,
-                kernel_size=15,
-                padding="same",
-            ),
-            nn.GELU(approximate="tanh"),
-        )
-        downsampled_seq_lens = [
-            self.cell_len // (2**i) for i in range(len(filter_list) - 1)
-        ]
-
-        self.conv_tower = nn.ModuleList(
-            [
-                ConvTowerBlock(
-                    dim_in=self._filter_list[i],
-                    dim_out=self._filter_list[i + 1],
-                    kernel_size=5,
-                    seq_len=seq_len,
-                    num_cells=config.num_cells,
-                )
-                for i, seq_len in zip(range(len(filter_list) - 1), downsampled_seq_lens)
-            ]
-        )
-
-        self.deconv_tower = nn.ModuleList(
-            [
-                DeConvTowerBlock(
-                    dim_in=filter_list[-1 - i],
-                    dim_out=filter_list[-1 - i - 1],
-                    kernel_size=5,
-                    stride=2,
-                    seq_len=seq_len // 2,
-                    num_cells=config.num_cells,
-                )
-                for i, seq_len in zip(
-                    range(len(filter_list) - 1), downsampled_seq_lens[::-1]
-                )
-            ]
-        )
-        self.transformer_layers = nn.ModuleList(
-            [
-                SelfAttentionBlock(
-                    num_heads=config.attention_heads,
-                    embed_dim=config.embed_dim,
-                    ffn_embed_dim=config.ffn_embed_dim,
-                    key_size=config.key_size,
-                    add_bias_kv=False,
-                    add_bias_fnn=False,
-                    ffn_activation_name="swish",
-                    use_glu_in_ffn=True,
-                    layer_norm_eps=1e-5,  # this is the default haiku value
-                    pre_layer_norm=True,
-                    name=f"attention_layer_{layer_idx}",
-                    rotary_embedding_config=self._rotary_embedding_config,
-                )
-                for layer_idx in range(config.num_layers)
-            ]
-        )
-
-        self.lm_head = LMHead(
-            dim_in=config.token_embed_dim,
-            embed_dim=config.embed_dim,
-            dim_out=config.alphabet_size,
-            num_hidden_layers=config.num_hidden_layers_head,
-        )
-
-    def forward(self, input_ids: torch.Tensor) -> dict[str, torch.Tensor]:
-        outs = {}
-        embeddings = self.token_embed(input_ids)
-        x = embeddings.permute(0, 2, 1)
-        x = self.stem_conv(x)
-        residuals = []
-        for _idx, conv_block in enumerate(self.conv_tower):
-            x, res = conv_block(x)
-            residuals.append(res)
-        residuals = residuals[::-1]
-        x = x.permute(0, 2, 1)
-
-        for layer_idx, transformer in enumerate(self.transformer_layers):
-            output = transformer(x)
-            x = output["embeddings"]
-            if (layer_idx + 1) in self.config.embeddings_layers_to_save:
-                outs[f"embeddings_{(layer_idx + 1)}"] = output["embeddings"]
-            if (layer_idx + 1) in self._attention_layers_to_save:
-                for map_number in self._attention_maps_per_layer_to_save[layer_idx + 1]:
-                    dkey = f"attention_map_layer_{layer_idx + 1}_number_{map_number}"
-                    outs[dkey] = output["attention_weights"][:, map_number + 1]
-        x = x.permute(0, 2, 1)
-        for deconv_block, res in zip(self.deconv_tower, residuals):
-            x = deconv_block(x, res)
-        x = x.permute(0, 2, 1)
-        logits = self.lm_head(x)
-        outs["logits"] = logits
-
-        return outs