Math error

app/image_captioning_service.py

@@ -6,6 +6,7 @@ import nltk
 import pickle
 import warnings
 import logging
+import math
 warnings.filterwarnings("ignore")
 
 # Configure logging
@@ -184,7 +185,6 @@ class EncoderCNN(torch.nn.Module):
 class PositionalEncoding(torch.nn.Module):
     def __init__(self, d_model, max_len=5000):
         super(PositionalEncoding, self).__init__()
-        import math
         
         # Create positional encoding
         pe = torch.zeros(max_len, d_model)
@@ -208,6 +208,9 @@ class TransformerDecoder(torch.nn.Module):
         super(TransformerDecoder, self).__init__()
         import math
         
+        # Store math module as an instance variable so we can use it in forward
+        self.math = math
+        
         # Embedding layer
         self.embedding = torch.nn.Embedding(vocab_size, embed_dim)
         self.positional_encoding = PositionalEncoding(embed_dim)
@@ -241,7 +244,7 @@ class TransformerDecoder(torch.nn.Module):
         tgt_mask = self.generate_square_subsequent_mask(tgt.size(1)).to(tgt.device)
         
         # Embed tokens and add positional encoding
-        tgt = self.embedding(tgt) * math.sqrt(self.embedding.embedding_dim)
+        tgt = self.embedding(tgt) * self.math.sqrt(self.embedding.embedding_dim)
         tgt = self.positional_encoding(tgt)
         tgt = self.dropout(tgt)
         
@@ -262,6 +265,9 @@ class ImageCaptioningModel(torch.nn.Module):
     def __init__(self, vocab_size, embed_dim, hidden_dim, num_heads, num_layers):
         super(ImageCaptioningModel, self).__init__()
         
+        # Make sure math is available
+        self.math = math
+        
         # Image encoder
         self.encoder = EncoderCNN(embed_dim)
         
@@ -295,7 +301,7 @@ class ImageCaptioningModel(torch.nn.Module):
             img_features = img_features.unsqueeze(1)
             
             # Start with < SOS > token
-            current_ids = torch.tensor([[vocab.word2idx['<SOS>']]], dtype=torch.long).to(image.device)
+            current_ids = torch.tensor([[vocab.word2idx['< SOS >']]], dtype=torch.long).to(image.device)
             
             # Generate words one by one
             result_caption = []