Create model.py

model.py

@@ -0,0 +1,134 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import lightning
+from safetensors.torch import save_file
+
+class Config:
+    vocab_size = 50304
+    n_epochs = 50
+    batch_size = 36
+    lr = 3e-4
+    wd = 1e-6
+    n_embed = 256
+    num_blocks = 12
+    num_heads = 12
+    head_size = n_embed//num_heads
+    context_len = 224
+    attn_dropout_val = 0.2
+    mha_dropout_val = 0.2
+    ffn_dropout_val = 0.2
+
+class CausalAttentionHead(nn.Module):
+    def __init__(self, config):
+        super(CausalAttentionHead, self).__init__()
+        self.config = config
+
+        self.query = nn.Linear(config.n_embed, config.head_size, bias=False)
+        self.key = nn.Linear(config.n_embed, config.head_size, bias=False)
+        self.value = nn.Linear(config.n_embed, config.head_size, bias=False)
+        self.attn_drop = nn.Dropout(config.attn_dropout_val)
+        # mask for causal attention during training
+        self.register_buffer("mask", torch.tril(torch.ones(config.context_len, config.context_len)))
+
+    def forward(self, x):
+        bs, context_len, embed_dim = x.shape
+        q, k, v = self.query(x), self.key(x), self.value(x)
+        attn_filter = torch.divide(torch.bmm(q, k.transpose(1, 2)), self.config.head_size)
+        attn_filter = attn_filter.masked_fill(self.mask[:context_len, :context_len]==0, float("-inf"))
+        attn_weights = F.softmax(attn_filter, dim=-1)
+        attn_weights = self.attn_drop(attn_weights)
+        output = torch.bmm(attn_weights, v)
+        return output
+
+class MultiHeadedAttention(nn.Module):
+    def __init__(self, config):
+        super(MultiHeadedAttention, self).__init__()
+        self.config = config
+        self.heads = nn.ModuleList(
+            [CausalAttentionHead(config) for _ in range(config.num_heads)]
+        )
+        self.proj = nn.Linear(config.num_heads*config.head_size, config.n_embed)
+        self.mha_drop = nn.Dropout(config.mha_dropout_val)
+
+    def forward(self, x):
+        mha_output = torch.cat([head(x) for head in self.heads], dim=-1)
+        return self.mha_drop(self.proj(mha_output))
+
+class FeedForwardNetwork(nn.Module):
+    def __init__(self, config):
+        super(FeedForwardNetwork, self).__init__()
+
+        self.ffn = nn.Sequential(
+            nn.Linear(config.n_embed, config.n_embed*4),
+            nn.GELU(),
+            nn.Linear(config.n_embed*4, config.n_embed),
+            nn.Dropout()
+        )
+    def forward(self, x):
+        return self.ffn(x)
+
+class Block(nn.Module):
+    def __init__(self, config):
+        super(Block, self).__init__()
+        self.mha = MultiHeadedAttention(config)
+        self.ln1 = nn.LayerNorm(config.n_embed)
+        self.ffn = FeedForwardNetwork(config)
+        self.ln2 = nn.LayerNorm(config.n_embed)
+
+    def forward(self, x):
+        x = self.ln1(x+self.mha(x))
+        x = self.ln2(x+self.ffn(x))
+        return x
+
+class GPT(lightning.LightningModule):
+    def __init__(self, config):
+        super(GPT, self).__init__()
+        self.config = config
+        self.save_hyperparameters()
+        self.token_embedding = nn.Embedding(config.vocab_size, config.n_embed)
+        self.positional_embedding = nn.Embedding(config.context_len, config.n_embed)
+        self.backbone = nn.Sequential(*[Block(config) for _ in range(config.num_blocks)])
+        self.lm_head = nn.Linear(config.n_embed, config.vocab_size)
+
+    def forward(self, x):
+        tok_emb = self.token_embedding(x)
+        pos_emb = self.positional_embedding(torch.arange(x.shape[1], device=self.device))
+        x = tok_emb+pos_emb
+        x = self.backbone(x)
+        logits = self.lm_head(x)
+        return logits
+
+    def get_loss(self, predictions, target):
+        B, C, V = predictions.shape
+        predictions = predictions.view(B*C, V)
+        target = target.view(B*C)
+        loss = F.cross_entropy(predictions, target)
+        return loss
+
+    def training_step(self, batch, batch_idx):
+        text, target = batch
+        text = text.long()
+        target = target.long()
+        logits = self(text)
+        loss = self.get_loss(logits, target)
+
+        self.log('loss', loss.item(), prog_bar=True)
+        logs = {'loss': loss}
+
+        return {"log": logs, "loss": loss}
+
+    def training_end(self, outputs):
+        avg_loss = torch.stack([x['log']['loss'] for x in outputs]).mean()
+        logs = {"log": avg_loss}
+        print(f"val_loss: {avg_loss}")
+        return {"log": logs}
+
+    def configure_optimizers(self):
+        opt = torch.optim.AdamW(self.parameters(), lr=self.config.lr, weight_decay=self.config.wd)
+        return [opt], []
+
+if __name__ == "__main__":
+    config = Config() 
+    gpt = GPT(config)
+    save_file(gpt, "storyGPT.safetensors")