Update evo_model.py

evo_model.py

@@ -1,8 +1,9 @@
+import os
+import json
 import torch
 import torch.nn as nn
 from transformers import PreTrainedModel, PretrainedConfig
 
-
 class EvoTransformerConfig(PretrainedConfig):
     def __init__(
         self,
@@ -30,7 +31,7 @@ class EvoTransformerForClassification(PreTrainedModel):
         super().__init__(config)
         self.config = config
 
-        # Expose architecture traits for dashboard or mutation
+        # === Architecture traits for UI, mutation, etc.
         self.num_layers = config.num_layers
         self.num_heads = config.num_heads
         self.ffn_dim = config.ffn_dim
@@ -42,7 +43,8 @@ class EvoTransformerForClassification(PreTrainedModel):
             nn.TransformerEncoderLayer(
                 d_model=config.hidden_size,
                 nhead=config.num_heads,
-                dim_feedforward=config.ffn_dim
+                dim_feedforward=config.ffn_dim,
+                batch_first=False  # Required for transpose trick
             )
             for _ in range(config.num_layers)
         ])
@@ -56,32 +58,37 @@ class EvoTransformerForClassification(PreTrainedModel):
         self.init_weights()
 
     def forward(self, input_ids, attention_mask=None, labels=None):
-        x = self.embedding(input_ids)  # [batch, seq_len, hidden_size]
-        x = x.transpose(0, 1)  # Transformer expects [seq_len, batch, hidden_size]
+        # Embedding and prep for transformer
+        x = self.embedding(input_ids)  # [batch, seq_len, hidden]
+        x = x.transpose(0, 1)          # [seq_len, batch, hidden]
+
+        key_padding_mask = (attention_mask == 0) if attention_mask is not None else None
 
         for layer in self.layers:
-            x = layer(x, src_key_padding_mask=(attention_mask == 0) if attention_mask is not None else None)
+            x = layer(x, src_key_padding_mask=key_padding_mask)
 
-        x = x.mean(dim=0)  # mean pooling
+        x = x.mean(dim=0)  # [batch, hidden] — mean pooling
         logits = self.classifier(x)
 
         if labels is not None:
             loss = nn.functional.cross_entropy(logits, labels)
             return loss, logits
+
         return logits
 
     def save_pretrained(self, save_directory):
-        import os, json
         os.makedirs(save_directory, exist_ok=True)
-        torch.save(self.state_dict(), f"{save_directory}/pytorch_model.bin")
-        with open(f"{save_directory}/config.json", "w") as f:
+        torch.save(self.state_dict(), os.path.join(save_directory, "pytorch_model.bin"))
+        with open(os.path.join(save_directory, "config.json"), "w") as f:
             f.write(self.config.to_json_string())
 
     @classmethod
     def from_pretrained(cls, load_directory):
-        config_path = f"{load_directory}/config.json"
-        model_path = f"{load_directory}/pytorch_model.bin"
+        config_path = os.path.join(load_directory, "config.json")
+        model_path = os.path.join(load_directory, "pytorch_model.bin")
         config = EvoTransformerConfig.from_json_file(config_path)
         model = cls(config)
-        model.load_state_dict(torch.load(model_path, map_location="cpu"))
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        model.load_state_dict(torch.load(model_path, map_location=device))
+        model.to(device)
         return model