Update modeling_lightning.py

modeling_lightning.py

@@ -5,7 +5,162 @@ import torch.nn as nn
 from transformers import PreTrainedModel, PretrainedConfig
 from transformers.modeling_outputs import CausalLMOutput
 
-from model import TransformerLanguageModel
+import torch
+import torch.nn.functional as F
+import torch.nn as nn
+import math
+
+embedding = 256
+heads = 4
+layers = 4
+dropout = 0.1
+msl = 160
+
+class PositionalEncoding(nn.Module):
+
+    def __init__(self, d_model, max_len=5000):
+
+        super(PositionalEncoding, self).__init__()
+
+        pe = torch.zeros(max_len, d_model)
+
+        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
+
+        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
+
+        pe[:, 0::2] = torch.sin(position * div_term)
+
+        pe[:, 1::2] = torch.cos(position * div_term)
+
+        pe = pe.unsqueeze(0).transpose(0, 1)
+
+        self.register_buffer('pe', pe)
+
+
+    def forward(self, x):
+
+        return x + self.pe[:x.size(1), :].transpose(0, 1)
+
+
+class CausalSelfAttention(nn.Module):
+    def __init__(self, d_model, nhead, dropout=0.1):
+        super().__init__()
+
+        assert d_model % nhead == 0
+
+        self.nhead = nhead
+        self.head_dim = d_model // nhead
+        self.dropout = dropout
+
+        self.qkv = nn.Linear(d_model, d_model * 3)
+        self.out_proj = nn.Linear(d_model, d_model)
+
+    def forward(self, x):
+        B, T, C = x.shape
+
+        # Create Q, K, V
+        q, k, v = self.qkv(x).chunk(3, dim=-1)
+
+        # [B, T, C] -> [B, heads, T, head_dim]
+        q = q.view(B, T, self.nhead, self.head_dim).transpose(1, 2)
+        k = k.view(B, T, self.nhead, self.head_dim).transpose(1, 2)
+        v = v.view(B, T, self.nhead, self.head_dim).transpose(1, 2)
+
+        y = F.scaled_dot_product_attention(
+            q,
+            k,
+            v,
+            attn_mask=None,
+            dropout_p=self.dropout if self.training else 0.0,
+            is_causal=True
+        )
+
+        y = y.transpose(1, 2).contiguous().view(B, T, C)
+
+        return self.out_proj(y)
+
+
+class TransformerBlock(nn.Module):
+    def __init__(self, d_model, nhead, dropout=0.1):
+        super().__init__()
+
+        self.norm1 = nn.LayerNorm(d_model)
+        self.attention = CausalSelfAttention(
+            d_model,
+            nhead,
+            dropout
+        )
+
+        self.norm2 = nn.LayerNorm(d_model)
+
+        self.ffn = nn.Sequential(
+            nn.Linear(d_model, d_model * 4),
+            nn.GELU(),
+            nn.Linear(d_model * 4, d_model),
+            nn.Dropout(dropout)
+        )
+
+    def forward(self, x):
+        x = x + self.attention(self.norm1(x))
+
+        x = x + self.ffn(self.norm2(x))
+
+        return x
+
+
+class TransformerLanguageModel(nn.Module):
+    def __init__(
+        self,
+        vocab_size,
+        d_model=512,
+        nhead=8,
+        num_layers=8,
+        dropout=0.1,
+        max_seq_len=160
+    ):
+        super().__init__()
+
+        self.d_model = d_model
+        self.max_seq_len = max_seq_len
+
+        self.token_embedding = nn.Embedding(
+            vocab_size,
+            d_model
+        )
+
+        self.positional_encoding = PositionalEncoding(
+            d_model,
+            max_seq_len
+        )
+
+        self.transformer = nn.ModuleList([
+            TransformerBlock(
+                d_model,
+                nhead,
+                dropout
+            )
+            for _ in range(num_layers)
+        ])
+
+        self.final_norm = nn.LayerNorm(d_model)
+        
+        self.output_layer = nn.Linear(
+        d_model,
+        vocab_size,
+        bias=False
+        )
+
+    def forward(self, src):
+        x = self.token_embedding(src)
+
+        x = self.positional_encoding(x)
+
+        for layer in self.transformer:
+            x = layer(x)
+
+        x = self.final_norm(x)
+
+        return self.output_layer(x)
 
 
 class LightningConfig(PretrainedConfig):