Update evo_model.py

evo_model.py

@@ -1,25 +1,23 @@
-# evo_model.py — EvoDecoderModel with fixed positional encoding (max_len=512)
 import torch
 import torch.nn as nn
+import math
 
-# Positional encoding used by transformer decoders
 class PositionalEncoding(nn.Module):
-    def __init__(self, d_model, max_len=512):  # Increased max_len
+    def __init__(self, d_model, max_len=512):
         super().__init__()
         pe = torch.zeros(max_len, d_model)
-        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
-        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-torch.log(torch.tensor(10000.0)) / d_model))
+        position = torch.arange(0, max_len, dtype=torch.float32).unsqueeze(1)
+        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
         pe[:, 0::2] = torch.sin(position * div_term)
         pe[:, 1::2] = torch.cos(position * div_term)
-        pe = pe.unsqueeze(0)  # shape: (1, max_len, d_model)
+        pe = pe.unsqueeze(0)  # [1, max_len, d_model]
         self.register_buffer('pe', pe)
 
     def forward(self, x):
-        return x + self.pe[:, :x.size(1)]  # Match sequence length
+        return x + self.pe[:, :x.size(1)]
 
-# Main EvoDecoder model
 class EvoDecoderModel(nn.Module):
-    def __init__(self, vocab_size, d_model=384, nhead=6, num_layers=6, dim_feedforward=1024, dropout=0.1):
+    def __init__(self, vocab_size, d_model=512, nhead=8, num_layers=6, dim_feedforward=2048, dropout=0.1):
         super().__init__()
         self.embedding = nn.Embedding(vocab_size, d_model)
         self.pos_encoder = PositionalEncoding(d_model)
@@ -36,12 +34,8 @@ class EvoDecoderModel(nn.Module):
     def forward(self, input_ids):
         embedded = self.embedding(input_ids)
         embedded = self.pos_encoder(embedded)
-
-        # Create causal mask for autoregressive decoding
         seq_len = embedded.size(1)
         mask = torch.triu(torch.ones(seq_len, seq_len, device=embedded.device), diagonal=1).bool()
-
-        # Use the input itself as memory for self-decoding
         output = self.decoder(embedded, embedded, tgt_mask=mask)
         logits = self.linear(output)
         return logits