Also figure out ways to improve my attention mechanism while figuring out if the multiple heads help or not. You cannot use the dot product as it is expensive to calculate.
Think how a separate parallel model could be used to guide the output of the first. Including ways to improve information density by maybe adding trainable memory parameters or some other way.

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import PreTrainedModel, PretrainedConfig

import torch
import torch.nn as nn
import torch.nn.functional as F

class CELU4(nn.Module):
    def __init__(self, alpha=1.0):
        super(CELU4, self).__init__()
        self.alpha = alpha  # Alpha parameter for CELU
    
    def forward(self, x):
        # Apply the CELU activation function
        celu_output = F.celu(x, alpha=self.alpha)
        # Clamp the output between 0 and 4
        celu6_output = torch.clamp(celu_output, max=4)
        return celu6_output
        
class Attention(nn.Module):
    def __init__(self, input_dim, num_heads):
        super(Attention, self).__init__()
        self.input_dim = input_dim
        self.num_heads = num_heads
        self.head_dim = input_dim // num_heads

        self.query_transform = nn.Linear(self.input_dim, self.input_dim)
        self.key_transform = nn.Linear(self.input_dim, self.input_dim)
        self.value_transform = nn.Linear(self.input_dim, self.input_dim)

        # Gating mechanism with learned sigmoid gate
        self.gate_transform = nn.Linear(self.head_dim, self.head_dim)
        self.celu = CELU4()

    def forward(self, input_seq):
        query = self.query_transform(input_seq)
        key = self.key_transform(input_seq)
        value = self.value_transform(input_seq)

        # Reshape into heads
        query = query.view(-1, input_seq.size(1), self.num_heads, self.head_dim).transpose(1, 2)
        key = key.view(-1, input_seq.size(1), self.num_heads, self.head_dim).transpose(1, 2)
        value = value.view(-1, input_seq.size(1), self.num_heads, self.head_dim).transpose(1, 2)

        # Without bilinear interaction, just multiply keys with gated attention
        gate = self.celu(self.gate_transform(query))
        attention_weights = gate * key
        attention_weights = F.softmax(attention_weights, dim=-1)

        # Combine heads
        output = attention_weights * value
        output = output.transpose(1, 2).contiguous().view(-1, input_seq.size(1), self.input_dim)

        return output



class TransformerLayer(nn.Module):
    def __init__(self, embed_size, num_heads, hidden_size):
        super(TransformerLayer, self).__init__()
        self.self_attention = Attention(embed_size, num_heads)
        self.feed_forward = nn.Sequential(
            nn.Linear(embed_size, hidden_size),
            nn.Linear(hidden_size, embed_size)
        )
        self.layer_norm1 = nn.LayerNorm(embed_size)
        self.layer_norm2 = nn.LayerNorm(embed_size)


    def forward(self, x):
        out = self.self_attention(x)
        out = self.layer_norm1(out + x)
        out = torch.relu(self.feed_forward(out))
        out = self.layer_norm2(out + out)
        return out


class LanguageModelConfig(PretrainedConfig):
    model_type = "gated_attention_transformer"

    def __init__(self, vocab_size=30522, embed_size=768, num_layers=6, num_heads=8, hidden_size=2048, **kwargs):
        super().__init__(**kwargs)
        self.vocab_size = vocab_size
        self.embed_size = embed_size
        self.num_layers = num_layers
        self.num_heads = num_heads
        self.hidden_size = hidden_size


class LanguageModel(PreTrainedModel):
    config_class = LanguageModelConfig

    def __init__(self, config):
        super().__init__(config)
        self.token_embedding = nn.Embedding(config.vocab_size, config.embed_size, padding_idx = 0)
        #self.position_encoding = nn.Embedding(4096, config.embed_size)
        self.transformer = nn.ModuleList(
            [TransformerLayer(config.embed_size, config.num_heads, config.hidden_size) for _ in range(config.num_layers)]
        )
        self.decoder = nn.Linear(config.embed_size, config.vocab_size)

        # Initialize weights
        self.init_weights()

    def forward(self, input_ids, attention_mask=None):
        batch_size = input_ids.size(0)
        seq_len = input_ids.size(1)
        #position_ids = torch.arange(seq_len, dtype=torch.long, device=input_ids.device).unsqueeze(0).repeat(batch_size, 1)

        token_embeddings = self.token_embedding(input_ids)
       # position_embeddings = self.position_encoding(position_ids)
        embeddings = token_embeddings #+ position_embeddings

        out = embeddings
        for layer in self.transformer:
            out = layer(out)

        out = self.decoder(out)
        return out

    def _init_weights(self, module):
        """Initialize the weights."""
        if isinstance(module, nn.Linear):
            module.weight.data.normal_(mean=0.0, std=0.02)
            if module.bias is not None:
                module.bias.data.zero_()
        elif isinstance(module, nn.Embedding):
            module.weight.data.normal_(mean=0.0, std=0.02)
        elif isinstance(module, nn.LayerNorm):
            module.bias.data.zero_()
            module.weight.data.fill_(1.0)