Upload MM_LLMs

config.json

@@ -1,5 +1,5 @@
 {
-  "_name_or_path": "multimodal-tinyllama-whisper-small-siglip/checkpoint-2500",
+  "_name_or_path": "multimodal-tinyllama-whisper-small-siglip/checkpoint-15500",
   "architectures": [
     "MM_LLMs"
   ],
@@ -206,6 +206,10 @@
   },
   "audio_conv_kernel": 240,
   "audio_conv_stride": 220,
+  "auto_map": {
+    "AutoConfig": "modeling.MM_LLMs_Config",
+    "AutoModel": "modeling.MM_LLMs"
+  },
   "hidden_size": 2048,
   "image_config": {
     "_name_or_path": "google/siglip-base-patch16-224",

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:9537a5b0cbb8c353c742959d43e527634e05a97ae06aed7ac23f668a1c06d924
-size 3509162622
+oid sha256:fd41c7ff4ba5d41aa0fb55a38f382e50d72f1da8d2521ff30c40d4e3a5a6bfb9
+size 3504434068

modeling.py

@@ -83,21 +83,16 @@ class MM_LLMs(PreTrainedModel):
         attn_dropout = 0.1
         is_add_bias_kv = True
         is_add_zero_attn = True
-        self.temporal_self_attention = nn.MultiheadAttention(
-            config.image_config.text_config.hidden_size,
-            config.attention_heads,
-            dropout=attn_dropout,
-            add_bias_kv=is_add_bias_kv,
-            add_zero_attn=is_add_zero_attn)
 
+        self.num_heads = config.attention_heads * 2
         self.audio_align_attention = nn.MultiheadAttention(config.llm_config.hidden_size,
-                                                           config.attention_heads * 2,
+                                                           self.num_heads,
                                                            dropout=attn_dropout,
                                                            add_bias_kv=is_add_bias_kv,
                                                            add_zero_attn=is_add_zero_attn)
 
         self.image_align_attention = nn.MultiheadAttention(config.llm_config.hidden_size,
-                                                           config.attention_heads * 2,
+                                                           self.num_heads,
                                                            dropout=attn_dropout,
                                                            add_bias_kv=is_add_bias_kv,
                                                            add_zero_attn=is_add_zero_attn)
@@ -123,12 +118,7 @@ class MM_LLMs(PreTrainedModel):
             self.config.image_config.text_config.hidden_size,
             bias=False)
 
-        self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
-
         self.layer_norm = nn.LayerNorm(config.image_config.text_config.hidden_size)
-        self.softmax = nn.Softmax(dim=-1)
-
-        self.sigmoid = nn.Sigmoid()
 
         self.loss_fct = CrossEntropyLoss()
 
@@ -188,9 +178,7 @@ class MM_LLMs(PreTrainedModel):
             images=None,
             audios=None,
             audio_starts=None,
-            audio_ends=None,
             image_starts=None,
-            image_ends=None,
             attention_mask=None,
             labels=None,
             audio_index=None,
@@ -212,7 +200,6 @@ class MM_LLMs(PreTrainedModel):
         if audio_features is not None:
 
             audio_starts = embed_tokens(audio_starts).unsqueeze(1)
-            audio_ends = embed_tokens(audio_ends).unsqueeze(1)
 
             audio_features = self.project_audio(
                 audio_features.transpose(
@@ -232,29 +219,39 @@ class MM_LLMs(PreTrainedModel):
                  dim),
                 device=token_embeddings.device,
                 dtype=token_embeddings.dtype)
+            new_audio_mask = torch.ones(
+                (
+                    token_embeddings.shape[1] * self.num_heads,
+                    seq_img * max_count,
+                    token_embeddings.shape[0]
+                ),
+                device=token_embeddings.device,
+                dtype=torch.bool)
             current_dim = 0
             for no, index in enumerate(audio_index):
                 if no > 0 and audio_index[no - 1] == index:
                     current_dim += 1
                 else:
                     current_dim = 0
-                new_audio[index, current_dim *
-                          seq_img: (current_dim + 1) * seq_img] = audio_features[no]
-                last_index = audio_index[0]
+                new_audio[
+                    index, current_dim *
+                    seq_img: (current_dim + 1) * seq_img
+                ] = audio_features[no]
+                new_audio_mask[index * self.num_heads: (index + 1) * self.num_heads, current_dim *
+                               seq_img: (current_dim + 1) * seq_img] = 0
 
             audio_features = self.audio_align_attention(
                 new_audio.transpose(
                     0,
                     1),
                 token_embeddings,
-                token_embeddings)[0].transpose(
+                token_embeddings,
+                attn_mask=new_audio_mask
+            )[0].transpose(
                 0,
                 1).contiguous()
 
-            # audio_features = add_positional_encoding(audio_features)
-
-            audio_inputs = torch.cat(
-                [torch.cat([audio_starts, audio_features], dim=1), audio_ends], dim=1)
+            audio_inputs = torch.cat([audio_starts, audio_features], dim=1)
 
             text_embeddings = torch.cat(
                 [torch.cat([text_embeddings[:, 0, :].unsqueeze(1), audio_inputs], dim=1), text_embeddings[:, 1:, :]],
@@ -265,7 +262,6 @@ class MM_LLMs(PreTrainedModel):
         if image_features is not None:
 
             image_starts = embed_tokens(image_starts).unsqueeze(1)
-            image_ends = embed_tokens(image_ends).unsqueeze(1)
 
             image_features = self.project_image(
                 image_features.transpose(
@@ -286,6 +282,16 @@ class MM_LLMs(PreTrainedModel):
                 device=token_embeddings.device,
                 dtype=token_embeddings.dtype)
 
+            new_img_mask = torch.ones(
+                (
+                    token_embeddings.shape[1] * self.num_heads,
+                    seq_img * max_count,
+                    token_embeddings.shape[0]
+                ),
+                device=token_embeddings.device,
+                dtype=torch.bool
+            )
+
             current_dim = 0
             for no, index in enumerate(image_index):
                 if no > 0 and image_index[no - 1] == index:
@@ -294,21 +300,21 @@ class MM_LLMs(PreTrainedModel):
                     current_dim = 0
                 new_img[index, current_dim *
                         seq_img: (current_dim + 1) * seq_img] = image_features[no]
-                last_index = image_index[0]
+                new_audio_mask[index * self.num_heads: (index + 1) * self.num_heads, current_dim *
+                               seq_img: (current_dim + 1) * seq_img] = 0
 
             image_features = self.image_align_attention(
                 new_img.transpose(
                     0,
                     1),
                 token_embeddings,
-                token_embeddings)[0].transpose(
+                token_embeddings,
+                attn_mask=new_img_mask,
+            )[0].transpose(
                 0,
                 1).contiguous()
 
-            # image_features = add_positional_encoding(image_features)
-
-            image_inputs = torch.cat(
-                [torch.cat([image_starts, image_features], dim=1), image_ends], dim=1)
+            image_inputs = torch.cat([image_starts, image_features], dim=1)
 
             text_embeddings = torch.cat(
                 [torch.cat([text_embeddings[:, 0, :].unsqueeze(1), image_inputs], dim=1),