Update text_image_audio.py

text_image_audio.py

@@ -83,18 +83,29 @@ class AudioEncoder(nn.Module):
         return self.forward(inputs)
 
 class ModalityTokenEncoder(nn.Module):
-    def __init__(self, projection_dim=CFG.projection_dim, token_size=CFG.token_size, device='cpu', *args, **kwargs):
+    def __init__(self, projection_dim=CFG.projection_dim, token_size=CFG.token_size, device='cpu', token_dim=CFG.token_dim, *args, **kwargs):
         super(ModalityTokenEncoder, self).__init__(*args, **kwargs)
         # Attributes
         self.projection_dim = projection_dim
         self.device = device
         self.token_size = token_size
+        self.token_dim = token_dim
         # Models
         audio_variance = torch.rand(1) * 0.5 + 0.1
         self.audio_token = nn.Parameter(torch.normal(mean=0, std=audio_variance.item(),
-                                                      size=(self.token_size, self.projection_dim)).to(self.device))
+                                                      size=(self.token_size, self.token_dim)).to(self.device))
+
+        self.token_projection = nn.Sequential(
+            nn.Linear(self.token_dim, 64),
+            nn.ReLU(),
+            nn.Linear(64, 128),
+            nn.ReLU(),
+            nn.Linear(128, self.projection_dim),
+            nn.LayerNorm(self.projection_dim)
+        )
+
     def forward(self):
-        return self.audio_token
+        return self.token_projection(self.audio_token)
 
     def __call__(self):
         return self.forward()
@@ -205,4 +216,4 @@ class OneEncoder(nn.Module, PyTorchModelHubMixin):
             #    fig.suptitle(display(Audio(query['input_values'], rate=self.sample_rate)))
             #plt.show()
         #return values, indices
-       
+