Update evo_model.py

evo_model.py

@@ -1,45 +1,30 @@
 import torch
 import torch.nn as nn
-import torch.nn.functional as F
-
-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)
-        memory_out = x[:, 0]
-        return self.memory_proj(memory_out)
+from torch.nn import TransformerEncoder, TransformerEncoderLayer
 
 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)
-
-    def forward(self, x):
-        x = self.encoder(x)
-        return self.classifier(x)
+    def __init__(self, vocab_size=30522, d_model=384, nhead=6, num_layers=6, dim_feedforward=1024, dropout=0.1, num_labels=2):
+        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.norm = nn.LayerNorm(d_model)
+
+        self.memory_proj = nn.Linear(d_model, d_model)
+        self.classifier = nn.Linear(d_model, num_labels)
+
+    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)
+
+        x = self.transformer(x)
+        x = self.norm(x)
+
+        memory_output = self.memory_proj(x[:, 0])
+        logits = self.classifier(memory_output)
+
+        return logits