Delete model

model/py

@@ -1,134 +0,0 @@
-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")