Update evo_model.py

evo_model.py

@@ -1,45 +1,50 @@
 import torch
 import torch.nn as nn
-import math
+import torch.nn.functional as F
 
-class MultiHeadSelfAttention(nn.Module):
-    def __init__(self, d_model=256, nhead=4):
+class SelfAttention(nn.Module):
+    def __init__(self, d_model, nhead):
         super().__init__()
-        self.nhead = nhead
-        self.d_head = d_model // nhead
-        self.qkv_proj = nn.Linear(d_model, d_model * 3)
+        self.qkv_proj = nn.Linear(d_model, 3 * d_model)
         self.out_proj = nn.Linear(d_model, d_model)
+        self.nhead = nhead
+        self.d_model = d_model
 
     def forward(self, x):
         B, T, C = x.size()
-        qkv = self.qkv_proj(x).view(B, T, self.nhead, 3 * self.d_head)
+        qkv = self.qkv_proj(x)
         q, k, v = qkv.chunk(3, dim=-1)
-        q, k, v = [t.transpose(1, 2) for t in (q, k, v)]
-        scores = (q @ k.transpose(-2, -1)) / math.sqrt(self.d_head)
+
+        q = q.view(B, T, self.nhead, C // self.nhead).transpose(1, 2)
+        k = k.view(B, T, self.nhead, C // self.nhead).transpose(1, 2)
+        v = v.view(B, T, self.nhead, C // self.nhead).transpose(1, 2)
+
+        scores = torch.matmul(q, k.transpose(-2, -1)) / (C // self.nhead) ** 0.5
         attn = torch.softmax(scores, dim=-1)
-        context = attn @ v
-        context = context.transpose(1, 2).contiguous().view(B, T, C)
-        return self.out_proj(context)
+        out = torch.matmul(attn, v)
+
+        out = out.transpose(1, 2).contiguous().view(B, T, C)
+        return self.out_proj(out)
 
 class FeedForward(nn.Module):
-    def __init__(self, d_model=256, d_ff=512):
+    def __init__(self, d_model, dim_feedforward):
         super().__init__()
         self.net = nn.Sequential(
-            nn.Linear(d_model, d_ff),
+            nn.Linear(d_model, dim_feedforward),
             nn.ReLU(),
-            nn.Linear(d_ff, d_model)
+            nn.Linear(dim_feedforward, d_model)
         )
 
     def forward(self, x):
         return self.net(x)
 
-class DecoderBlock(nn.Module):
-    def __init__(self, d_model=256, nhead=4, d_ff=512):
+class TransformerBlock(nn.Module):
+    def __init__(self, d_model, nhead, dim_feedforward):
         super().__init__()
+        self.attn = SelfAttention(d_model, nhead)
         self.ln1 = nn.LayerNorm(d_model)
-        self.attn = MultiHeadSelfAttention(d_model, nhead)
+        self.ffn = FeedForward(d_model, dim_feedforward)
         self.ln2 = nn.LayerNorm(d_model)
-        self.ffn = FeedForward(d_model, d_ff)
 
     def forward(self, x):
         x = x + self.attn(self.ln1(x))
@@ -47,22 +52,21 @@ class DecoderBlock(nn.Module):
         return x
 
 class EvoDecoder(nn.Module):
-    def __init__(self, vocab_size=50257, d_model=256, nhead=4, num_layers=3, d_ff=512, max_len=512):
+    def __init__(self, vocab_size, d_model=256, nhead=4, num_layers=3, dim_feedforward=1024):
         super().__init__()
         self.token_emb = nn.Embedding(vocab_size, d_model)
-        self.pos_emb = nn.Embedding(max_len, d_model)
-        self.blocks = nn.ModuleList([
-            DecoderBlock(d_model, nhead, d_ff) for _ in range(num_layers)
+        self.pos_emb = nn.Embedding(512, d_model)
+        self.blocks = nn.Sequential(*[
+            TransformerBlock(d_model, nhead, dim_feedforward) for _ in range(num_layers)
         ])
         self.ln_f = nn.LayerNorm(d_model)
         self.fc_out = nn.Linear(d_model, vocab_size)
 
     def forward(self, x):
-        B, T = x.shape
-        token = self.token_emb(x)
-        pos = self.pos_emb(torch.arange(T, device=x.device).unsqueeze(0))
-        x = token + pos
-        for block in self.blocks:
-            x = block(x)
+        B, T = x.size()
+        tok = self.token_emb(x)
+        pos = self.pos_emb(torch.arange(T, device=x.device).unsqueeze(0).expand(B, T))
+        x = tok + pos
+        x = self.blocks(x)
         x = self.ln_f(x)
         return self.fc_out(x)