Update evo_model.py

evo_model.py

@@ -1,22 +1,68 @@
-import torch
 import torch.nn as nn
+import torch
+
+class FeedForward(nn.Module):
+    def __init__(self, dim, hidden_dim):
+        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),
+        )
+
+    def forward(self, x):
+        return self.net(x)
+
+class Attention(nn.Module):
+    def __init__(self, dim, heads=4):
+        super().__init__()
+        self.heads = heads
+        self.scale = dim ** -0.5
+
+        self.qkv_proj = nn.Linear(dim, dim * 3)
+        self.out_proj = nn.Linear(dim, dim)
+
+    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)
+
+class TransformerBlock(nn.Module):
+    def __init__(self, dim, heads, hidden_dim):
+        super().__init__()
+        self.attn = Attention(dim, heads)
+        self.ffn = FeedForward(dim, hidden_dim)
+        self.ln1 = nn.LayerNorm(dim)
+        self.ln2 = nn.LayerNorm(dim)
+
+    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, d_model=256, nhead=4, num_layers=3, dim_feedforward=1024, dropout=0.1):
-        super(EvoDecoderModel, self).__init__()
-        self.embedding = nn.Embedding(vocab_size, d_model)
-        self.pos_embedding = nn.Parameter(torch.zeros(1, 512, d_model))  # max length 512
-        decoder_layer = nn.TransformerDecoderLayer(d_model, nhead, dim_feedforward, dropout)
-        self.transformer_decoder = nn.TransformerDecoder(decoder_layer, num_layers)
-        self.output_layer = nn.Linear(d_model, vocab_size)
-
-    def forward(self, tgt, memory=None):
-        seq_len = tgt.size(1)
-        embedded = self.embedding(tgt) + self.pos_embedding[:, :seq_len, :]
-
-        # If no memory is provided, use dummy memory filled with zeros
-        if memory is None:
-            memory = torch.zeros_like(embedded)
-
-        output = self.transformer_decoder(embedded.transpose(0, 1), memory.transpose(0, 1))
-        return self.output_layer(output.transpose(0, 1))
+    def __init__(self, vocab_size, dim=256, depth=3, heads=4, hidden_dim=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)
+
+    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)
+        x = self.ln_f(x)
+        logits = self.fc_out(x)
+        return logits