import torch.nn as nn
from src.model.norm import LayerNorm
from src.model.feed_forward import FeedForwardNetwork
from src.model.attention import MultiHeadAttention
from src.model.config import GPTConfig


class Transformer(nn.Module):
    def __init__(self, config: GPTConfig):
        super().__init__()
        self.norm1 = LayerNorm(config.embed_dim)
        self.attn = MultiHeadAttention(
            config.embed_dim, config.head_dim, config.drop_rate, config.num_heads
        )
        self.norm2 = LayerNorm(config.embed_dim)
        self.ff = FeedForwardNetwork(config.embed_dim)
        self.drop = nn.Dropout(config.drop_rate)

    def forward(self, x):
        x = x + self.drop(self.attn(self.norm1(x)))
        x = x + self.drop(self.ff(self.norm2(x)))
        return x