Update evo_model.py
Browse files- evo_model.py +3 -5
evo_model.py
CHANGED
@@ -4,26 +4,24 @@ import torch.nn as nn
|
|
4 |
import math
|
5 |
|
6 |
class PositionalEncoding(nn.Module):
|
7 |
-
def __init__(self, d_model, max_len=
|
8 |
super().__init__()
|
9 |
pe = torch.zeros(max_len, d_model)
|
10 |
position = torch.arange(0, max_len).unsqueeze(1)
|
11 |
div_term = torch.exp(torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model))
|
12 |
pe[:, 0::2] = torch.sin(position * div_term)
|
13 |
pe[:, 1::2] = torch.cos(position * div_term)
|
14 |
-
pe = pe.unsqueeze(0) #
|
15 |
self.register_buffer('pe', pe)
|
16 |
|
17 |
def forward(self, x):
|
18 |
-
if x.size(1) > self.pe.size(1):
|
19 |
-
raise ValueError(f"Input sequence length {x.size(1)} exceeds max_len {self.pe.size(1)}")
|
20 |
return x + self.pe[:, :x.size(1)]
|
21 |
|
22 |
class EvoDecoderModel(nn.Module):
|
23 |
def __init__(self, vocab_size, d_model=512, nhead=8, num_layers=6, dim_feedforward=2048, dropout=0.1):
|
24 |
super().__init__()
|
25 |
self.token_embed = nn.Embedding(vocab_size, d_model)
|
26 |
-
self.pos_encoder = PositionalEncoding(d_model)
|
27 |
decoder_layer = nn.TransformerDecoderLayer(
|
28 |
d_model=d_model,
|
29 |
nhead=nhead,
|
|
|
4 |
import math
|
5 |
|
6 |
class PositionalEncoding(nn.Module):
|
7 |
+
def __init__(self, d_model, max_len=128): # ❗ Reverted to 128 to match saved model
|
8 |
super().__init__()
|
9 |
pe = torch.zeros(max_len, d_model)
|
10 |
position = torch.arange(0, max_len).unsqueeze(1)
|
11 |
div_term = torch.exp(torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model))
|
12 |
pe[:, 0::2] = torch.sin(position * div_term)
|
13 |
pe[:, 1::2] = torch.cos(position * div_term)
|
14 |
+
pe = pe.unsqueeze(0) # (1, max_len, d_model)
|
15 |
self.register_buffer('pe', pe)
|
16 |
|
17 |
def forward(self, x):
|
|
|
|
|
18 |
return x + self.pe[:, :x.size(1)]
|
19 |
|
20 |
class EvoDecoderModel(nn.Module):
|
21 |
def __init__(self, vocab_size, d_model=512, nhead=8, num_layers=6, dim_feedforward=2048, dropout=0.1):
|
22 |
super().__init__()
|
23 |
self.token_embed = nn.Embedding(vocab_size, d_model)
|
24 |
+
self.pos_encoder = PositionalEncoding(d_model, max_len=128) # ❗ Match checkpoint shape
|
25 |
decoder_layer = nn.TransformerDecoderLayer(
|
26 |
d_model=d_model,
|
27 |
nhead=nhead,
|