Transformer Day: Annotated miniGPT about lxmls-toolkit HOT 9 OPEN

class Block(nn.Module):
    """ an unassuming Transformer block """

    def __init__(self, config):
        super().__init__()
        self.ln_1 = nn.LayerNorm(config.n_embd)
        self.attn = CausalSelfAttention(config)
        self.ln_2 = nn.LayerNorm(config.n_embd)

        self.mlp = nn.Sequential(
            nn.Linear(config.n_embd, 4 * config.n_embd),
            NewGELU(),
            nn.Linear(4 * config.n_embd, config.n_embd),
            nn.Dropout(config.resid_pdrop)
        )

    def forward(self, x):
        block_in = x
        # x = x + self.attn(self.ln_1(x))
        attn_out = block_in + self.attn(self.ln_1(block_in))

        # x = x + self.mlp(self.ln_2(x))
        block_out = attn_out + self.mlp(self.ln_2(attn_out))
        return block_out

class Block(nn.Module):
    """ an unassuming Transformer block """

    def __init__(self, config):
        super().__init__()
        self.ln_1 = nn.LayerNorm(config.n_embd)
        self.attn = CausalSelfAttention(config)
        self.ln_2 = nn.LayerNorm(config.n_embd)
        self.mlp = nn.ModuleDict(dict(
            c_fc    = nn.Linear(config.n_embd, 4 * config.n_embd),
            c_proj  = nn.Linear(4 * config.n_embd, config.n_embd),
            act     = NewGELU(),
            dropout = nn.Dropout(config.resid_pdrop),
        ))
        m = self.mlp
        self.mlpf = lambda x: m.dropout(m.c_proj(m.act(m.c_fc(x)))) # MLP forward

    def forward(self, x):
        x = x + self.attn(self.ln_1(x))
        x = x + self.mlpf(self.ln_2(x))
        return x

class Block(nn.Module):
    """ an unassuming Transformer block """

    def __init__(self, config):
        super().__init__()
        self.layer_norm_1 = nn.LayerNorm(config.n_embd)
        self.dropout_1 = nn.Dropout(config.resid_pdrop)
        self.causal_multi_head_attention = CausalSelfAttention(config)
        self.layer_norm_2 = nn.LayerNorm(config.n_embd)
        self.dropout_2 = nn.Dropout(config.att_pdrop)
        self.mlp = nn.ModuleDict(dict(
            c_fc    = nn.Linear(config.n_embd, 4 * config.n_embd),
            c_proj  = nn.Linear(4 * config.n_embd, config.n_embd),
            act     = NewGELU()
        ))

    def feed_forward(self, x):
        return self.mlp.c_proj(self.mlp.act(self.mlp.c_fc(x)))

    def forward(self, x):
        # NOTE: We need to remove dropout from CausalSelfAttention
        x = x + self.dropout_1(self.causal_multi_head_attention(self.layer_norm_1(x)))
        x = x + self.dropout_2(self.feed_forward(self.layer_norm_2(x)))
        return x

class Block(nn.Module):
    """ an unassuming Transformer block """

    def __init__(self, config):
        super().__init__()

        # Causal Multi Head Attention
        # NOTE: We need to remove dropout from CausalSelfAttention
        self.causal_multi_head_attention = CausalSelfAttention(config)

        # Feeed Forward
        self.feed_forward = nn.Sequential(
            nn.Linear(config.n_embd, 4 * config.n_embd),
            NewGELU(),
            nn.Linear(4 * config.n_embd, config.n_embd),
        )

        # Dropout and Layer normalization
        self.layer_norm_1 = nn.LayerNorm(config.n_embd)
        self.dropout_1 = nn.Dropout(config.resid_pdrop)
        self.layer_norm_2 = nn.LayerNorm(config.n_embd)
        self.dropout_2 = nn.Dropout(config.att_pdrop)

    def forward(self, x):
        block_in = x
        attn_out = block_in + self.dropout_1(self.causal_multi_head_attention(self.layer_norm_1(block_in)))
        block_out = attn_out + self.dropout_2(self.feed_forward(self.layer_norm_2(attn_out)))
        return block_out