Update evo_model.py

evo_model.py

@@ -3,7 +3,7 @@ import torch.nn as nn
 import math
 
 class MultiHeadSelfAttention(nn.Module):
-    def __init__(self, d_model, nhead):
+    def __init__(self, d_model=256, nhead=4):
         super().__init__()
         self.nhead = nhead
         self.d_head = d_model // nhead
@@ -15,7 +15,6 @@ class MultiHeadSelfAttention(nn.Module):
         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)]
-
         scores = (q @ k.transpose(-2, -1)) / math.sqrt(self.d_head)
         attn = torch.softmax(scores, dim=-1)
         context = attn @ v
@@ -23,7 +22,7 @@ class MultiHeadSelfAttention(nn.Module):
         return self.out_proj(context)
 
 class FeedForward(nn.Module):
-    def __init__(self, d_model, d_ff):
+    def __init__(self, d_model=256, d_ff=512):
         super().__init__()
         self.net = nn.Sequential(
             nn.Linear(d_model, d_ff),
@@ -35,7 +34,7 @@ class FeedForward(nn.Module):
         return self.net(x)
 
 class DecoderBlock(nn.Module):
-    def __init__(self, d_model, nhead, d_ff):
+    def __init__(self, d_model=256, nhead=4, d_ff=512):
         super().__init__()
         self.ln1 = nn.LayerNorm(d_model)
         self.attn = MultiHeadSelfAttention(d_model, nhead)
@@ -48,7 +47,7 @@ class DecoderBlock(nn.Module):
         return x
 
 class EvoDecoder(nn.Module):
-    def __init__(self, vocab_size, d_model=512, nhead=8, num_layers=6, d_ff=2048, max_len=512):
+    def __init__(self, vocab_size=50257, d_model=256, nhead=4, num_layers=3, d_ff=512, max_len=512):
         super().__init__()
         self.token_emb = nn.Embedding(vocab_size, d_model)
         self.pos_emb = nn.Embedding(max_len, d_model)