fixes

code/models/transformer.py

@@ -0,0 +1,40 @@
+import math
+
+import torch
+from torch import nn
+
+
+class TransformerModel(nn.Module):
+    def __init__(self, input_dim: int,
+                 output_dim: int,
+                 seq_len: int,
+                 embed_dim: int,
+                 num_heads: int,
+                 num_enc_layers: int):
+        super().__init__()
+        self.input_proj = nn.Linear(input_dim, embed_dim)
+
+        # Compute positional embedding ONCE at init
+        pe = self._get_sinusoidal_embedding(seq_len, embed_dim)  # (seq_len, embed_dim)
+        self.register_buffer('pos_embed', pe.unsqueeze(0))  # (1, seq_len, embed_dim)
+
+        encoder_layer = nn.TransformerEncoderLayer(embed_dim, num_heads)
+        self.encoder = nn.TransformerEncoder(encoder_layer, num_enc_layers)
+        self.pool = nn.AdaptiveAvgPool1d(1)
+        self.head = nn.Linear(embed_dim, output_dim)
+
+    def _get_sinusoidal_embedding(self, seq_len, embed_dim):
+        position = torch.arange(0, seq_len).unsqueeze(1)
+        div_term = torch.exp(torch.arange(0, embed_dim, 2) * -(math.log(10000.0) / embed_dim))
+        pe = torch.zeros(seq_len, embed_dim)
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
+        return pe  # (seq_len, embed_dim)
+
+    def forward(self, x):
+        # x: (B, seq_len, input_dim)
+        x = self.input_proj(x) + self.pos_embed[:, :x.size(1), :]  # broadcasting
+        x = x.permute(1, 0, 2)  # (S, B, E)
+        enc = self.encoder(x)  # (S, B, E)
+        pooled = enc.mean(0)  # (B, E)
+        return self.head(pooled)