Create modeling.py

modeling.py

@@ -0,0 +1,73 @@
+from transformers import ASTModel, ViTModel, PretrainedConfig, PreTrainedModel
+import numpy as np
+import torch
+import torch.nn as nn
+from einops import reduce
+
+
+class MuVis(nn.Module):
+  def __init__(self, embed_dims=768, latent_dims=128, sampling_rate=16000):
+    super(MuVis, self).__init__()
+    self.sampling_rate = sampling_rate
+    self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
+
+    self.ast = ASTModel.from_pretrained("MIT/ast-finetuned-audioset-10-10-0.4593", low_cpu_mem_usage=True)
+    self.wav_lin = nn.Linear(embed_dims, latent_dims)
+
+    self.vit = ViTModel.from_pretrained("google/vit-base-patch16-224-in21k", low_cpu_mem_usage=True)
+    self.img_lin = nn.Linear(embed_dims, latent_dims)
+
+
+  def forward(self, wav=None, img=None):
+    wav_out = None
+    img_out = None
+
+    if wav is not None:
+      wav_out = self.ast(**wav)["last_hidden_state"]
+      wav_out = self.wav_lin(wav_out)
+      wav_out = reduce(wav_out, "b n d -> b d", "mean")
+      wav_out = wav_out / wav_out.norm(dim=-1, keepdim=True)
+
+    if img is not None:
+      img_out = self.vit(**img)["last_hidden_state"]
+      img_out = self.img_lin(img_out)
+      img_out = reduce(img_out, "b n d -> b d", "mean")
+      img_out = img_out / img_out.norm(dim=-1, keepdim=True)
+
+
+    assert wav_out is not None or img_out is not None
+
+    if wav_out is None or img_out is None:
+      return wav_out if img_out is None else img_out
+    return (wav_out, img_out)
+
+
+class MuVisConfig(PretrainedConfig):
+    model_type = "muvis"
+
+    def __init__(
+        self,
+        embed_dims=768,
+        latent_dims=128,
+        sampling_rate=16000,
+        **kwargs,
+    ):
+      self.embed_dims = embed_dims
+      self.latent_dims = latent_dims
+      self.sampling_rate = sampling_rate
+      super().__init__(**kwargs)
+
+
+class MuVisModel(PreTrainedModel):
+    config_class = MuVisConfig
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = MuVis(
+            embed_dims=config.embed_dims,
+            latent_dims=config.latent_dims,
+            sampling_rate=config.sampling_rate,
+        )
+
+    def forward(self, wav=None, img=None):
+        return self.model(wav=wav, img=img)