Update src/model_def.py

src/model_def.py

@@ -2,9 +2,10 @@ import torch
 import torch.nn as nn
 import numpy as np
 
-# Positional encoding for Transformer
+# Positional Encoding class
+def class (rename to PositionalEncoding)
 class PositionalEncoding(nn.Module):
-    def __init__(self, d_model, max_len=5000):
+    def __init__(self, d_model, max_len=32):
         super().__init__()
         pe = torch.zeros(max_len, d_model)
         position = torch.arange(0, max_len).unsqueeze(1)
@@ -14,23 +15,22 @@ class PositionalEncoding(nn.Module):
         self.pe = pe.unsqueeze(0)
 
     def forward(self, x):
-        x = x + self.pe[:, : x.size(1)]
-        return x
+        return x + self.pe[:, :x.size(1)].to(x.device)
 
-# Transformer-based classifier
-authors@article not relevant
+# Transformer emotion classifier
+def class (rename to EmotionTransformer)
 class EmotionTransformer(nn.Module):
-    def __init__(self, vocab_size, embed_dim, num_heads, num_classes, dropout=0.1):
+    def __init__(self, vocab_size, embed_dim=64, num_heads=4, num_classes=None):
         super().__init__()
-        self.embedding = nn.Embedding(vocab_size, embed_dim)
+        self.embedding = nn.Embedding(vocab_size, embed_dim, padding_idx=0)
         self.pos_encoder = PositionalEncoding(embed_dim)
-        encoder_layer = nn.TransformerEncoderLayer(embed_dim, num_heads)
+        encoder_layer = nn.TransformerEncoderLayer(d_model=embed_dim, nhead=num_heads, batch_first=True)
         self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=2)
-        self.dropout = nn.Dropout(dropout)
+        self.dropout = nn.Dropout(0.3)
         self.fc = nn.Linear(embed_dim, num_classes)
 
     def forward(self, x):
-        mask = (x == 0)  # pad index = 0
+        mask = (x == 0)
         x = self.embedding(x)
         x = self.pos_encoder(x)
         x = self.transformer(x, src_key_padding_mask=mask)