Update evo_model.py

evo_model.py

@@ -1,26 +1,45 @@
 import torch
 import torch.nn as nn
-from torch.nn import TransformerEncoder, TransformerEncoderLayer
+import torch.nn.functional as F
 
-class EvoTransformer(nn.Module):
-    def __init__(self, vocab_size=30522, d_model=512, nhead=8, num_layers=6, dim_feedforward=1024, dropout=0.1):
-        super(EvoTransformer, self).__init__()
-        self.embedding = nn.Embedding(vocab_size, d_model)
-        self.memory_token = nn.Parameter(torch.zeros(1, 1, d_model))
-        
-        encoder_layer = TransformerEncoderLayer(d_model=d_model, nhead=nhead, dim_feedforward=dim_feedforward, dropout=dropout)
-        self.transformer = TransformerEncoder(encoder_layer, num_layers=num_layers)
-        
-        self.memory_proj = nn.Linear(d_model, d_model)
-        # 🔥 self.norm removed
-
-    def forward(self, input_ids):
-        x = self.embedding(input_ids)
-        memory_token = self.memory_token.expand(x.size(0), -1, -1)
-        x = torch.cat([memory_token, x], dim=1)
+class TransformerEncoder(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.embedding = nn.Embedding(config["vocab_size"], config["d_model"])
+        encoder_layer = nn.TransformerEncoderLayer(
+            d_model=config["d_model"],
+            nhead=config["nhead"],
+            dim_feedforward=config["ff_dim"],
+            dropout=0.1,
+            activation="gelu",
+            batch_first=True,
+        )
+        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=config["num_layers"])
+        self.memory_token = nn.Parameter(torch.randn(1, 1, config["d_model"]))
+        self.memory_proj = nn.Linear(config["d_model"], config["d_model"])
 
+    def forward(self, x):
+        x = self.embedding(x)
+        B, T, D = x.shape
+        memory = self.memory_token.repeat(B, 1, 1)
+        x = torch.cat([memory, x], dim=1)
         x = self.transformer(x)
-        # 🔥 x = self.norm(x) removed
+        memory_out = x[:, 0]
+        return self.memory_proj(memory_out)
+
+class EvoTransformer(nn.Module):
+    def __init__(self):
+        super().__init__()
+        config = {
+            "vocab_size": 30522,
+            "d_model": 384,
+            "nhead": 6,
+            "ff_dim": 1024,
+            "num_layers": 6,
+        }
+        self.encoder = TransformerEncoder(config)
+        self.classifier = nn.Linear(config["d_model"], 2)
 
-        memory_output = self.memory_proj(x[:, 0])
-        return memory_output
+    def forward(self, x):
+        x = self.encoder(x)
+        return self.classifier(x)