Update evo_model.py

evo_model.py

@@ -1,68 +1,69 @@
-import torch.nn as nn
 import torch
+import torch.nn as nn
+import math
 
-class FeedForward(nn.Module):
-    def __init__(self, dim, hidden_dim):
+class MultiHeadSelfAttention(nn.Module):
+    def __init__(self, d_model, nhead):
         super().__init__()
-        self.net = nn.Sequential(
-            nn.Linear(dim, hidden_dim),
-            nn.GELU(),
-            nn.Dropout(0.1),
-            nn.Linear(hidden_dim, dim),
-            nn.Dropout(0.1),
-        )
+        self.nhead = nhead
+        self.d_head = d_model // nhead
+        self.qkv_proj = nn.Linear(d_model, d_model * 3)
+        self.out_proj = nn.Linear(d_model, d_model)
 
     def forward(self, x):
-        return self.net(x)
+        B, T, C = x.size()
+        qkv = self.qkv_proj(x).view(B, T, self.nhead, 3 * self.d_head)
+        q, k, v = qkv.chunk(3, dim=-1)
+        q, k, v = [t.transpose(1, 2) for t in (q, k, v)]
 
-class Attention(nn.Module):
-    def __init__(self, dim, heads=4):
-        super().__init__()
-        self.heads = heads
-        self.scale = dim ** -0.5
+        scores = (q @ k.transpose(-2, -1)) / math.sqrt(self.d_head)
+        attn = torch.softmax(scores, dim=-1)
+        context = attn @ v
+        context = context.transpose(1, 2).contiguous().view(B, T, C)
+        return self.out_proj(context)
 
-        self.qkv_proj = nn.Linear(dim, dim * 3)
-        self.out_proj = nn.Linear(dim, dim)
+class FeedForward(nn.Module):
+    def __init__(self, d_model, d_ff):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(d_model, d_ff),
+            nn.ReLU(),
+            nn.Linear(d_ff, d_model)
+        )
 
     def forward(self, x):
-        B, T, C = x.shape
-        qkv = self.qkv_proj(x).reshape(B, T, 3, self.heads, C // self.heads).permute(2, 0, 3, 1, 4)
-        q, k, v = qkv[0], qkv[1], qkv[2]
-        attn_scores = (q @ k.transpose(-2, -1)) * self.scale
-        attn_weights = attn_scores.softmax(dim=-1)
-        attn_output = attn_weights @ v
-        attn_output = attn_output.transpose(1, 2).reshape(B, T, C)
-        return self.out_proj(attn_output)
+        return self.net(x)
 
-class TransformerBlock(nn.Module):
-    def __init__(self, dim, heads, hidden_dim):
+class DecoderBlock(nn.Module):
+    def __init__(self, d_model, nhead, d_ff):
         super().__init__()
-        self.attn = Attention(dim, heads)
-        self.ffn = FeedForward(dim, hidden_dim)
-        self.ln1 = nn.LayerNorm(dim)
-        self.ln2 = nn.LayerNorm(dim)
+        self.ln1 = nn.LayerNorm(d_model)
+        self.attn = MultiHeadSelfAttention(d_model, nhead)
+        self.ln2 = nn.LayerNorm(d_model)
+        self.ffn = FeedForward(d_model, d_ff)
 
     def forward(self, x):
         x = x + self.attn(self.ln1(x))
         x = x + self.ffn(self.ln2(x))
         return x
 
-class EvoDecoderModel(nn.Module):
-    def __init__(self, vocab_size, dim=256, depth=3, heads=4, hidden_dim=512):
+class EvoDecoder(nn.Module):
+    def __init__(self, vocab_size, d_model=512, nhead=8, num_layers=6, d_ff=2048, max_len=512):
         super().__init__()
-        self.token_emb = nn.Embedding(vocab_size, dim)
-        self.pos_emb = nn.Embedding(512, dim)
-        self.blocks = nn.Sequential(*[TransformerBlock(dim, heads, hidden_dim) for _ in range(depth)])
-        self.ln_f = nn.LayerNorm(dim)
-        self.fc_out = nn.Linear(dim, vocab_size)
+        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.ln_f = nn.LayerNorm(d_model)
+        self.fc_out = nn.Linear(d_model, vocab_size)
 
     def forward(self, x):
         B, T = x.shape
-        pos = torch.arange(0, T, device=x.device).unsqueeze(0)
-        tok = self.token_emb(x)
-        pos = self.pos_emb(pos)
-        x = tok + pos
-        x = self.blocks(x)
+        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)
         x = self.ln_f(x)
-        logits = self.fc_out(x)
-        return logits
+        return self.fc_out(x)