import torch
import torch.nn as nn

from .MuiltHeadAttention import MuiltHeadAttention_v1
from .FeedForward import FeedForward
from .LayerNormal import LayerNormal

class TransformersBlocks(nn.Module):
    """ 
    Treanforners模块 """
    def __init__(self,config):
        super().__init__()
        self.attention = MuiltHeadAttention_v1(
            d_in=config['emb_dim'],
            d_out=config['emb_dim'],
            num_tokens=config['context_length'],
            dropout_p=config['drop_rate'],
            num_head=config['n_heads'],
        )
        self.ff = FeedForward(config)

        self.norm1 = LayerNormal(config['emb_dim'])
        self.norm2 = LayerNormal(config['emb_dim'])

        self.drop = nn.Dropout(config['drop_rate'])
    
    def forward(self,x):
        raw_x = x
        # 第一层归一化
        x = self.norm1(x)
        # 获取上下文向量
        x = self.attention(x)
        x = self.drop(x)
        raw_x, x = x, x + raw_x

        # 第二层归一化
        x = self.norm2(x)
        x = self.ff(x) # 包含线性层、GRLU()、线性层
        x = self.drop(x)
        x = x + raw_x
        return x
