evo_model.py

# evo_model.py — EvoDecoderModel with fixed positional encoding (max_len=512)
import torch
import torch.nn as nn

# Positional encoding used by transformer decoders
class PositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=512):  # Increased max_len
        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))
        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)
        self.register_buffer('pe', pe)

    def forward(self, x):
        return x + self.pe[:, :x.size(1)]  # Match sequence length

# 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):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, d_model)
        self.pos_encoder = PositionalEncoding(d_model)
        decoder_layer = nn.TransformerDecoderLayer(
            d_model=d_model,
            nhead=nhead,
            dim_feedforward=dim_feedforward,
            dropout=dropout,
            batch_first=True
        )
        self.decoder = nn.TransformerDecoder(decoder_layer, num_layers=num_layers)
        self.linear = nn.Linear(d_model, vocab_size)

    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