rxnn 0.1.0__py3-none-any.whl

src/training/tokenizer.py

@@ -0,0 +1,208 @@
+import os
+from pathlib import Path
+from tokenizers import Tokenizer
+from tokenizers.models import BPE, WordPiece, Unigram, WordLevel
+from tokenizers.trainers import BpeTrainer, WordPieceTrainer, UnigramTrainer, WordLevelTrainer
+from tokenizers.pre_tokenizers import Whitespace, Punctuation, BertPreTokenizer, ByteLevel
+from tokenizers.processors import TemplateProcessing
+from tokenizers.normalizers import Lowercase, NFKC, Sequence
+from transformers import PreTrainedTokenizerFast
+from typing import Any
+
+class TokenizerTrainer:
+    def __init__(
+            self,
+            vocab_size: int = 30000,
+            model_type: str = "byte-level-bpe",  # Options: "bpe", "wordpiece", "unigram", "sentencepiece"
+            special_tokens: list[str] = None,
+            lowercase: bool = False,
+            normalization: bool = False,
+            pre_tokenizer_type: str = "bert",  # Options: "bert", "whitespace_punctuation",
+            vocab: Any = None,
+            byte_fallback: bool = False,
+            max_input_chars_per_word: int = 32,
+            use_post_processor: bool = True,
+            post_processor_single: str = "[BOS] $A [EOS]",
+            post_processor_pair: str = "[BOS] $A [EOS][BOS] $B:1 [EOS]:1",
+            post_processor_special_tokens: list[str] = None,
+    ):
+        self.vocab_size = vocab_size
+        self.special_tokens = special_tokens if special_tokens is not None else ["[PAD]", "[UNK]", "[BOS]", "[EOS]",
+                                                                                 "[MASK]"]
+        self.model_type = model_type.lower()
+        self.lowercase = lowercase
+        self.normalization = normalization
+        self.pre_tokenizer_type = pre_tokenizer_type.lower()
+
+        # Initialize tokenizer model
+        if self.model_type == "bpe":
+            self.tokenizer = Tokenizer(BPE(unk_token="[UNK]", vocab=vocab, byte_fallback=byte_fallback))
+        elif self.model_type == "wordpiece":
+            self.tokenizer = Tokenizer(
+                WordPiece(unk_token="[UNK]", vocab=vocab, max_input_chars_per_word=max_input_chars_per_word))
+        elif self.model_type == "unigram":
+            self.tokenizer = Tokenizer(Unigram(unk_id="[UNK]", vocab=None, byte_fallback=byte_fallback))
+        elif self.model_type == "wordlevel":
+            self.tokenizer = Tokenizer(WordLevel(unk_token="[UNK]", vocab=None))
+        elif self.model_type == "byte-level-bpe":
+            self.tokenizer = Tokenizer(BPE(unk_token="[UNK]", vocab=vocab, byte_fallback=byte_fallback))
+            self.tokenizer.pre_tokenizer = ByteLevel(add_prefix_space=True)
+        else:
+            raise ValueError(f"Unsupported model type: {model_type}")
+
+        # Configure pre-tokenizer
+        if self.model_type != "byte-level-bpe":
+            if self.pre_tokenizer_type == "bert":
+                self.tokenizer.pre_tokenizer = BertPreTokenizer()
+            elif self.pre_tokenizer_type == "whitespace_punctuation":
+                self.tokenizer.pre_tokenizer = Whitespace()
+                self.tokenizer.pre_tokenizer = Punctuation()
+            elif self.pre_tokenizer_type == "whitespace":
+                self.tokenizer.pre_tokenizer = Whitespace()
+            else:
+                raise ValueError(f"Unsupported pre-tokenizer: {pre_tokenizer_type}")
+
+        # Add normalization steps
+        if self.normalization:
+            normalizers = []
+            if self.lowercase:
+                normalizers.append(Lowercase())
+            normalizers.append(NFKC())
+            self.tokenizer.normalizer = Sequence(normalizers)
+
+        self.use_post_processor = use_post_processor
+        self.post_processor_single = post_processor_single
+        self.post_processor_pair = post_processor_pair
+        self.post_processor_special_tokens = post_processor_special_tokens
+
+    def train(
+            self,
+            files: list[str],
+            limit_alphabet: int = 1000,
+            show_progress: bool = True,
+            **kwargs
+    ):
+        # Prepare trainer based on model type
+        trainer_kwargs = {
+            "vocab_size": self.vocab_size,
+            "special_tokens": self.special_tokens,
+            "limit_alphabet": limit_alphabet,
+            "show_progress": show_progress,
+            **kwargs  # Allow custom parameters
+        }
+
+        if self.model_type in ["bpe", "byte-level-bpe"]:
+            trainer = BpeTrainer(**trainer_kwargs)
+        elif self.model_type == "wordpiece":
+            trainer = WordPieceTrainer(**trainer_kwargs)
+        elif self.model_type == "unigram":
+            trainer = UnigramTrainer(**trainer_kwargs)
+        elif self.model_type == "wordlevel":
+            trainer = WordLevelTrainer(**trainer_kwargs)
+
+        # Train tokenizer
+        self.tokenizer.train(files, trainer)
+
+        if self.use_post_processor:
+            post_processor_special_tokens = self.post_processor_special_tokens or ["[BOS]", "[EOS]"]
+            self.tokenizer.post_processor = TemplateProcessing(
+                single=self.post_processor_single,
+                pair=self.post_processor_pair,
+                special_tokens=[(token, self.tokenizer.token_to_id(token)) for token in post_processor_special_tokens],
+            )
+
+    def save(self, output_dir: str):
+        os.makedirs(output_dir, exist_ok=True)
+        self.tokenizer.save(f"{output_dir}/tokenizer.json")
+
+    def load(self, model_path: str):
+        self.tokenizer = Tokenizer.from_file(model_path)
+
+    def get_hf_tokenizer(self):
+        return PreTrainedTokenizerFast(
+            tokenizer_object=self.tokenizer,
+            unk_token="[UNK]",
+            pad_token="[PAD]",
+            cls_token="[CLS]",
+            sep_token="[SEP]",
+            mask_token="[MASK]"
+        )
+
+    def push_to_hub(
+            self,
+            repo_id: str,
+            create: bool = False,
+            private: bool = False,
+            api_token: str = None,
+            **kwargs
+    ):
+        """
+        Push the trained tokenizer to HuggingFace Hub.
+
+        Args:
+            repo_id (str): Hub repository name (e.g., "username/my-tokenizer")
+            private (bool): Whether the repo is private
+            api_token (str): HuggingFace API token (optional if already logged in)
+            **kwargs: Additional args for HuggingFace API
+        """
+        from huggingface_hub import HfApi, Repository
+
+        # Create a temporary directory for Hub upload
+        temp_dir = "temp_hub_upload"
+        os.makedirs(temp_dir, exist_ok=True)
+        self.save(temp_dir)  # Save tokenizer files locally
+
+        # Push to Hub using HuggingFace API
+        api = HfApi(token=api_token)
+        if create:
+            api.create_repo(
+                repo_id=repo_id,
+                private=private,
+                exist_ok=True,
+            )
+
+        # Push files to the repo
+        api.upload_folder(
+            repo_id=repo_id,
+            folder_path=temp_dir,
+            repo_type="model",
+            **kwargs
+        )
+
+        # Cleanup
+        os.remove(Path(temp_dir) / 'tokenizer.json')
+        os.rmdir(temp_dir)
+
+    @staticmethod
+    def hf_tokenizer_from_file(path: str):
+        return PreTrainedTokenizerFast(
+            tokenizer_file=path,
+            unk_token="[UNK]",
+            pad_token="[PAD]",
+            cls_token="[CLS]",
+            sep_token="[SEP]",
+            mask_token="[MASK]"
+        )
+
+    @classmethod
+    def from_pretrained(cls, repo_id: str, **kwargs):
+        """
+        Load tokenizer from HuggingFace Hub.
+
+        Args:
+            repo_id (str): Hub repository name (e.g., "username/my-tokenizer")
+            **kwargs: Additional args for HuggingFace API
+        """
+        from huggingface_hub import hf_hub_download
+
+        # Download tokenizer.json from Hub
+        tokenizer_file = hf_hub_download(
+            repo_id=repo_id,
+            filename="tokenizer.json",
+            **kwargs
+        )
+
+        # Initialize trainer and load tokenizer
+        trainer = cls()
+        trainer.load(tokenizer_file)
+        return trainer

src/transformers/attention.py

@@ -0,0 +1,324 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import math
+from positional import RotaryPositionalEmbedding, RelativePositionalEmbedding
+
+
+class MultiHeadAttention(nn.Module):
+    """Custom, extendable Multi-head attention layer, with RoPE support"""
+
+    def __init__(
+            self,
+            embed_dim: int,
+            num_heads: int,
+            dropout: float = 0.0,
+            rope: RotaryPositionalEmbedding = None,
+            rope_only_for_query: bool = False,
+            use_relative_embeddings: bool = False,
+            max_seq_len: int = 1024,
+            use_flash_attention: bool = False,
+            is_causal: bool = False,
+            use_bias: bool = False,
+            *args,
+            **kwargs,
+    ):
+        super(MultiHeadAttention, self).__init__(*args, **kwargs)
+        assert embed_dim % num_heads == 0, "embed_dim must be divisible by num_heads"
+        self.embed_dim = embed_dim
+        self.num_heads = num_heads
+
+        self.use_flash_attention = use_flash_attention
+        self.is_causal = is_causal
+        self.use_bias = use_bias
+        if use_relative_embeddings:
+            self.use_flash_attention = False
+            self.rel_embed = RelativePositionalEmbedding(max_seq_len, embed_dim // num_heads)
+            self.rope = None
+            self.rope_only_for_query = False
+        else:
+            self.rel_embed = None
+            self.rope = rope
+            self.rope_only_for_query = rope_only_for_query
+        self.dropout = nn.Dropout(dropout)
+        self._init_q(embed_dim)
+        self._init_kv(embed_dim)
+        self._init_out(embed_dim)
+
+    def _init_q(self, embed_dim: int):
+        """Initialize query projection"""
+        self.q_proj = nn.Linear(embed_dim, embed_dim, bias=self.use_bias)
+
+    def _init_kv(self, embed_dim: int):
+        """Initialize key and value projections"""
+        self.k_proj = nn.Linear(embed_dim, embed_dim, bias=self.use_bias)
+        self.v_proj = nn.Linear(embed_dim, embed_dim, bias=self.use_bias)
+
+    def _init_out(self, embed_dim: int):
+        """Initialize output projection"""
+        self.out_proj = nn.Linear(embed_dim, embed_dim)
+
+    def _forward_qkv(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, b: int, t: int, d: int):
+        """Forward pass through query, key, and value projections, and split the results into heads"""
+        q = self.q_proj(query).view(b, t, self.num_heads, d // self.num_heads).transpose(1, 2)
+        k = self.k_proj(key).view(b, -1, self.num_heads, d // self.num_heads).transpose(1, 2)
+        v = self.v_proj(value).view(b, -1, self.num_heads, d // self.num_heads).transpose(1, 2)
+        return q, k, v
+
+    def _apply_rope(self, q: torch.Tensor, k: torch.Tensor):
+        if self.rope is not None:
+            if self.rope_only_for_query:
+                q = self.rope.forward_one(q)
+            else:
+                q, k = self.rope(q, k)
+        return q, k
+
+    def _calculate_attn_weights(self, q: torch.Tensor, k: torch.Tensor, d: int, mask: torch.Tensor = None):
+        """Calculate attention weights using scaled dot-product attention"""
+        q, k = self._apply_rope(q, k)
+        attn_logits = torch.matmul(q, k.transpose(-2, -1)) / (d // self.num_heads) ** 0.5
+        if mask is not None:
+            attn_logits = attn_logits.masked_fill(mask == 0, float('-inf'))
+        return F.softmax(attn_logits, dim=-1)
+
+    def _calculate_attn_weight_with_relative_embeddings(self, q: torch.Tensor, k: torch.Tensor,
+                                                        mask: torch.Tensor = None):
+        """Calculate attention weights using scaled dot-product attention and apply relative embedding"""
+        attn_logits = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(q.size(-1))
+        rel_pos_bias = self.rel_embed(q, k)
+        attn_logits += rel_pos_bias
+        if mask is not None:
+            attn_logits = attn_logits.masked_fill(mask == 0, float('-inf'))
+        return F.softmax(attn_logits, dim=-1)
+
+    def _calculate_output(self, attn_weights: torch.Tensor, v: torch.Tensor, b: int, t: int, d: int):
+        """Calculate the output by multiplying attention weights with values and concatenating heads"""
+        return torch.matmul(attn_weights, v).transpose(1, 2).contiguous().view(b, t, d)
+
+    def _flash_attention(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, b: int, t: int, d: int,
+                         mask: torch.Tensor = None, enable_gqa: bool = False):
+        attn_output = F.scaled_dot_product_attention(
+            q, k, v,
+            attn_mask=mask if not self.is_causal else None,
+            dropout_p=self.dropout.p if self.training else 0.0,
+            is_causal=self.is_causal,
+            enable_gqa=enable_gqa,
+        )
+
+        # Reshape back to (B, T, D)
+        return attn_output.transpose(1, 2).contiguous().view(b, t, d)
+
+    def _calculate_flash_attention(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, b: int, t: int, d: int,
+                                   mask: torch.Tensor = None):
+        # Compute attention with FlashAttention
+        return self._flash_attention(q.contiguous(), k.contiguous(), v.contiguous(), b, t, d, mask=mask)
+
+    def forward(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, mask: torch.Tensor = None):
+        b, t, d = query.size()
+        q, k, v = self._forward_qkv(query, key, value, b, t, d)
+        if self.use_flash_attention:
+            q, k = self._apply_rope(q, k)
+            attn_output = self._calculate_flash_attention(q, k, v, b, t, d, mask=mask)
+        else:
+            if not self.rel_embed:
+                attn_weights = self._calculate_attn_weights(q, k, d, mask=mask)
+            else:
+                attn_weights = self._calculate_attn_weight_with_relative_embeddings(q, k, mask=mask)
+
+            attn_weights = self.dropout(attn_weights)
+
+            attn_output = self._calculate_output(attn_weights, v, b, t, d)
+        return self.out_proj(attn_output)
+
+
+class GroupedQueryAttention(MultiHeadAttention):
+    """Custom Grouped Query attention layer, with RoPE support"""
+
+    def __init__(
+            self,
+            embed_dim: int,
+            num_heads: int,
+            num_groups: int,
+            dropout: float = 0.0,
+            rope: RotaryPositionalEmbedding = None,
+            rope_only_for_query: bool = False,
+            use_relative_embeddings: bool = False,
+            max_seq_len: int = 1024,
+            use_flash_attention: bool = False,
+            is_causal: bool = False,
+            use_bias: bool = False,
+            *args,
+            **kwargs,
+    ):
+        self.num_groups = num_groups
+        super(GroupedQueryAttention, self).__init__(
+            embed_dim,
+            num_heads,
+            dropout=dropout,
+            rope=rope,
+            rope_only_for_query=rope_only_for_query,
+            use_relative_embeddings=use_relative_embeddings,
+            max_seq_len=max_seq_len,
+            use_flash_attention=use_flash_attention,
+            is_causal=is_causal,
+            use_bias=use_bias,
+            *args,
+            **kwargs,
+        )
+        assert num_heads % num_groups == 0, "num_heads must be divisible by num_groups"
+
+    def _init_kv(self, embed_dim: int):
+        self.k_proj = nn.Linear(embed_dim, embed_dim // (self.num_heads // self.num_groups), bias=self.use_bias)
+        self.v_proj = nn.Linear(embed_dim, embed_dim // (self.num_heads // self.num_groups), bias=self.use_bias)
+
+    def _forward_qkv(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, b: int, t: int, d: int):
+        """Override query, key, and value projections for GQA case - split data into heads and groups"""
+        head_dim = d // self.num_heads
+        if self.use_flash_attention:
+            q = self.q_proj(query).view(b, t, self.num_heads, head_dim).transpose(1, 2)
+            k = self.k_proj(key).view(b, -1, self.num_groups, head_dim).transpose(1, 2)
+            v = self.v_proj(value).view(b, -1, self.num_groups, head_dim).transpose(1, 2)
+        else:
+            group_heads = self.num_heads // self.num_groups
+
+            # Process Q
+            q = self.q_proj(query).view(b, t, self.num_groups, group_heads, head_dim).permute(0, 2, 3, 1,
+                                                                                              4)  # (B, G, group_heads, T, head_dim)
+
+            # Process K and V
+            k = self.k_proj(key).view(b, -1, self.num_groups, head_dim).transpose(1, 2)  # (B, G, S, head_dim)
+            v = self.v_proj(value).view(b, -1, self.num_groups, head_dim).transpose(1, 2)  # (B, G, S, head_dim)
+
+            # Expand and flatten to 4D tensors
+            k = k.unsqueeze(2).expand(-1, -1, group_heads, -1, -1)  # (B, G, group_heads, S, head_dim)
+            v = v.unsqueeze(2).expand(-1, -1, group_heads, -1, -1)  # (B, G, group_heads, S, head_dim)
+
+            q = q.flatten(start_dim=1, end_dim=2)  # (B, H, T, head_dim)
+            k = k.flatten(start_dim=1, end_dim=2)  # (B, H, S, head_dim)
+            v = v.flatten(start_dim=1, end_dim=2)  # (B, H, S, head_dim)
+        return q, k, v
+
+    def _calculate_flash_attention(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, b: int, t: int, d: int,
+                                   mask: torch.Tensor = None):
+        return self._flash_attention(
+            q.contiguous(), k.contiguous(), v.contiguous(), b, t, d, mask=mask,
+            enable_gqa=(self.num_heads != self.num_groups)
+        )
+
+
+class MultiQueryAttention(MultiHeadAttention):
+    """Custom Multi Query attention layer, with RoPE support"""
+
+    def __init__(
+            self,
+            embed_dim: int,
+            num_heads: int,
+            dropout: float = 0.0,
+            rope: RotaryPositionalEmbedding = None,
+            rope_only_for_query: bool = False,
+            use_relative_embeddings: bool = False,
+            max_seq_len: int = 1024,
+            use_flash_attention: bool = False,
+            is_causal: bool = False,
+            use_bias: bool = False,
+            *args,
+            **kwargs,
+    ):
+        super(MultiQueryAttention, self).__init__(
+            embed_dim,
+            num_heads,
+            dropout=dropout,
+            rope=rope,
+            rope_only_for_query=rope_only_for_query,
+            use_relative_embeddings=use_relative_embeddings,
+            max_seq_len=max_seq_len,
+            use_flash_attention=use_flash_attention,
+            is_causal=is_causal,
+            use_bias=use_bias,
+            *args,
+            **kwargs
+        )
+
+    def _init_kv(self, embed_dim: int):
+        """Override key/value initialization for MQA case"""
+        self.k_proj = nn.Linear(embed_dim, embed_dim // self.num_heads, bias=self.use_bias)
+        self.v_proj = nn.Linear(embed_dim, embed_dim // self.num_heads, bias=self.use_bias)
+
+    def _forward_qkv(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, b: int, t: int, d: int):
+        """Override query, key, and value projections for GQA case - use multiple heads
+        for query and single for key/values"""
+        if self.use_flash_attention:
+            q = self.q_proj(query).view(b, t, self.num_heads, d // self.num_heads).transpose(1, 2)
+            k = self.k_proj(key).view(b, -1, 1, d // self.num_heads).transpose(1, 2)
+            v = self.v_proj(value).view(b, -1, 1, d // self.num_heads).transpose(1, 2)
+        else:
+            q = self.q_proj(query).view(b, t, self.num_heads, d // self.num_heads).transpose(1, 2)
+            k = self.k_proj(key).unsqueeze(1).expand(-1, self.num_heads, -1, -1)
+            v = self.v_proj(value).unsqueeze(1).expand(-1, self.num_heads, -1, -1)
+        return q, k, v
+
+    def _calculate_flash_attention(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, b: int, t: int, d: int,
+                                   mask: torch.Tensor = None):
+        return self._flash_attention(
+            q.contiguous(), k.contiguous(), v.contiguous(), b, t, d, mask=mask,
+            enable_gqa=True
+        )
+
+
+def init_attention(
+        embed_dim: int,
+        num_heads: int,
+        attention_type: str,
+        gqa_groups: int = 1,
+        dropout: float = 0.0,
+        rope: RotaryPositionalEmbedding = None,
+        rope_only_for_query: bool = False,
+        use_relative_embeddings: bool = False,
+        max_seq_len: int = 1024,
+        use_flash_attention: bool = False,
+        is_causal: bool = False,
+        use_bias: bool = False,
+) -> MultiHeadAttention:
+    assert attention_type == 'mha' or attention_type == 'gqa' or attention_type == 'mqa', \
+        "Error, attention type should be one of: 'mha', 'gqa', 'mqa'"
+
+    if attention_type == "gqa":
+        return GroupedQueryAttention(
+            embed_dim,
+            num_heads,
+            gqa_groups,
+            dropout=dropout,
+            rope=rope,
+            use_relative_embeddings=use_relative_embeddings,
+            max_seq_len=max_seq_len,
+            rope_only_for_query=rope_only_for_query,
+            use_flash_attention=use_flash_attention,
+            is_causal=is_causal,
+            use_bias=use_bias,
+        )
+    elif attention_type == "mqa":
+        return MultiQueryAttention(
+            embed_dim,
+            num_heads,
+            dropout=dropout,
+            rope=rope,
+            use_relative_embeddings=use_relative_embeddings,
+            max_seq_len=max_seq_len,
+            rope_only_for_query=rope_only_for_query,
+            use_flash_attention=use_flash_attention,
+            is_causal=is_causal,
+            use_bias=use_bias,
+        )
+    else:
+        return MultiHeadAttention(
+            embed_dim,
+            num_heads,
+            dropout=dropout,
+            rope=rope,
+            use_relative_embeddings=use_relative_embeddings,
+            max_seq_len=max_seq_len,
+            rope_only_for_query=rope_only_for_query,
+            use_flash_attention=use_flash_attention,
+            is_causal=is_causal,
+            use_bias=use_bias,
+        )

src/transformers/ff.py

@@ -0,0 +1,72 @@
+import torch
+import torch.nn as nn
+
+
+class FeedForward(nn.Module):
+    """Basic Feed-forward layer with activation function and optional dropout"""
+
+    def __init__(self, embed_dim: int, hidden_dim: int, activation: nn.Module, dropout: float = 0.0, *args, **kwargs):
+        super(FeedForward, self).__init__(*args, **kwargs)
+        self.fc1 = nn.Linear(embed_dim, hidden_dim)
+        self.activation = activation
+        self.fc2 = nn.Linear(hidden_dim, embed_dim)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.fc1(x)
+        x = self.activation(x)
+        x = self.dropout(x)
+        return self.fc2(x)
+
+
+class GatedLinearUnit(nn.Module):
+    """Gated linear unit layer with configurable activation (SwiGLU, ReGLU, etc.)"""
+
+    def __init__(self, embed_dim: int, hidden_dim: int, activation: nn.Module, *args, **kwargs):
+        super(GatedLinearUnit, self).__init__(*args, **kwargs)
+        self.linear = nn.Linear(embed_dim, hidden_dim * 2)
+        self.activation = activation
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        l, g = self.linear(x).chunk(2, dim=-1)
+        return l * self.activation(g)
+
+
+class GatedFeedForward(nn.Module):
+    """Gated feed-forward layer with activation function and optional dropout"""
+
+    def __init__(self, embed_dim: int, hidden_dim: int, activation: nn.Module, dropout: float = 0.0, *args, **kwargs):
+        super(GatedFeedForward, self).__init__(*args, **kwargs)
+        self.fc1 = GatedLinearUnit(embed_dim, hidden_dim, activation)
+        self.fc2 = nn.Linear(hidden_dim, embed_dim)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.fc1(x)
+        x = self.dropout(x)
+        return self.fc2(x)
+
+
+class LinearActivation(nn.Module):
+    """Linear activation - identity function, for Bilinear Gated Unit"""
+
+    def __init__(self, *args, **kwargs):
+        super(LinearActivation, self).__init__(*args, **kwargs)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return x
+
+
+def get_activation_layer(activation: str):
+    if activation == 'relu':
+        return nn.ReLU()
+    elif activation == 'gelu':
+        return nn.GELU()
+    elif activation == 'silu' or activation == 'swish':
+        return nn.SiLU()
+    elif activation == 'sigmoid':
+        return nn.Sigmoid()
+    elif activation == 'linear':
+        return LinearActivation()
+    else:
+        raise ValueError(f'Activation {activation} not supported')