HemanM commited on
Commit
cf057f8
·
verified ·
1 Parent(s): 95f96da

Update evo_model.py

Browse files
Files changed (1) hide show
  1. 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=512): # Increased from 128 to 512
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) # Shape: (1, max_len, d_model)
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,