Upload MM_LLMs

README.md

@@ -0,0 +1,201 @@
+---
+library_name: transformers
+tags: []
+---
+
+# Model Card for Model ID
+
+<!-- Provide a quick summary of what the model is/does. -->
+
+
+
+## Model Details
+
+### Model Description
+
+<!-- Provide a longer summary of what this model is. -->
+
+This is the model card of a 🤗 transformers model that has been pushed on the Hub. This model card has been automatically generated.
+
+- **Developed by:** [More Information Needed]
+- **Funded by [optional]:** [More Information Needed]
+- **Shared by [optional]:** [More Information Needed]
+- **Model type:** [More Information Needed]
+- **Language(s) (NLP):** [More Information Needed]
+- **License:** [More Information Needed]
+- **Finetuned from model [optional]:** [More Information Needed]
+
+### Model Sources [optional]
+
+<!-- Provide the basic links for the model. -->
+
+- **Repository:** [More Information Needed]
+- **Paper [optional]:** [More Information Needed]
+- **Demo [optional]:** [More Information Needed]
+
+## Uses
+
+<!-- Address questions around how the model is intended to be used, including the foreseeable users of the model and those affected by the model. -->
+
+### Direct Use
+
+<!-- This section is for the model use without fine-tuning or plugging into a larger ecosystem/app. -->
+
+[More Information Needed]
+
+### Downstream Use [optional]
+
+<!-- This section is for the model use when fine-tuned for a task, or when plugged into a larger ecosystem/app -->
+
+[More Information Needed]
+
+### Out-of-Scope Use
+
+<!-- This section addresses misuse, malicious use, and uses that the model will not work well for. -->
+
+[More Information Needed]
+
+## Bias, Risks, and Limitations
+
+<!-- This section is meant to convey both technical and sociotechnical limitations. -->
+
+[More Information Needed]
+
+### Recommendations
+
+<!-- This section is meant to convey recommendations with respect to the bias, risk, and technical limitations. -->
+
+Users (both direct and downstream) should be made aware of the risks, biases and limitations of the model. More information needed for further recommendations.
+
+## How to Get Started with the Model
+
+Use the code below to get started with the model.
+
+[More Information Needed]
+
+## Training Details
+
+### Training Data
+
+<!-- This should link to a Dataset Card, perhaps with a short stub of information on what the training data is all about as well as documentation related to data pre-processing or additional filtering. -->
+
+[More Information Needed]
+
+### Training Procedure 
+
+<!-- This relates heavily to the Technical Specifications. Content here should link to that section when it is relevant to the training procedure. -->
+
+#### Preprocessing [optional]
+
+[More Information Needed]
+
+
+#### Training Hyperparameters
+
+- **Training regime:** [More Information Needed] <!--fp32, fp16 mixed precision, bf16 mixed precision, bf16 non-mixed precision, fp16 non-mixed precision, fp8 mixed precision -->
+
+#### Speeds, Sizes, Times [optional]
+
+<!-- This section provides information about throughput, start/end time, checkpoint size if relevant, etc. -->
+
+[More Information Needed]
+
+## Evaluation
+
+<!-- This section describes the evaluation protocols and provides the results. -->
+
+### Testing Data, Factors & Metrics
+
+#### Testing Data
+
+<!-- This should link to a Dataset Card if possible. -->
+
+[More Information Needed]
+
+#### Factors
+
+<!-- These are the things the evaluation is disaggregating by, e.g., subpopulations or domains. -->
+
+[More Information Needed]
+
+#### Metrics
+
+<!-- These are the evaluation metrics being used, ideally with a description of why. -->
+
+[More Information Needed]
+
+### Results
+
+[More Information Needed]
+
+#### Summary
+
+
+
+## Model Examination [optional]
+
+<!-- Relevant interpretability work for the model goes here -->
+
+[More Information Needed]
+
+## Environmental Impact
+
+<!-- Total emissions (in grams of CO2eq) and additional considerations, such as electricity usage, go here. Edit the suggested text below accordingly -->
+
+Carbon emissions can be estimated using the [Machine Learning Impact calculator](https://mlco2.github.io/impact#compute) presented in [Lacoste et al. (2019)](https://arxiv.org/abs/1910.09700).
+
+- **Hardware Type:** [More Information Needed]
+- **Hours used:** [More Information Needed]
+- **Cloud Provider:** [More Information Needed]
+- **Compute Region:** [More Information Needed]
+- **Carbon Emitted:** [More Information Needed]
+
+## Technical Specifications [optional]
+
+### Model Architecture and Objective
+
+[More Information Needed]
+
+### Compute Infrastructure
+
+[More Information Needed]
+
+#### Hardware
+
+[More Information Needed]
+
+#### Software
+
+[More Information Needed]
+
+## Citation [optional]
+
+<!-- If there is a paper or blog post introducing the model, the APA and Bibtex information for that should go in this section. -->
+
+**BibTeX:**
+
+[More Information Needed]
+
+**APA:**
+
+[More Information Needed]
+
+## Glossary [optional]
+
+<!-- If relevant, include terms and calculations in this section that can help readers understand the model or model card. -->
+
+[More Information Needed]
+
+## More Information [optional]
+
+[More Information Needed]
+
+## Model Card Authors [optional]
+
+[More Information Needed]
+
+## Model Card Contact
+
+[More Information Needed]
+
+

config.json

@@ -0,0 +1,505 @@
+{
+  "_name_or_path": "multimodal-mistral/checkpoint-1000",
+  "architectures": [
+    "MM_LLMs"
+  ],
+  "audio_config": {
+    "_name_or_path": "mesolitica/malaysian-whisper-small",
+    "activation_dropout": 0.0,
+    "activation_function": "gelu",
+    "add_cross_attention": false,
+    "apply_spec_augment": false,
+    "architectures": [
+      "WhisperForConditionalGeneration"
+    ],
+    "attention_dropout": 0.0,
+    "bad_words_ids": null,
+    "begin_suppress_tokens": [
+      220,
+      50257
+    ],
+    "bos_token_id": 50257,
+    "chunk_size_feed_forward": 0,
+    "classifier_proj_size": 256,
+    "cross_attention_hidden_size": null,
+    "d_model": 768,
+    "decoder_attention_heads": 12,
+    "decoder_ffn_dim": 3072,
+    "decoder_layerdrop": 0.0,
+    "decoder_layers": 12,
+    "decoder_start_token_id": 50258,
+    "diversity_penalty": 0.0,
+    "do_sample": false,
+    "dropout": 0.0,
+    "early_stopping": false,
+    "encoder_attention_heads": 12,
+    "encoder_ffn_dim": 3072,
+    "encoder_layerdrop": 0.0,
+    "encoder_layers": 12,
+    "encoder_no_repeat_ngram_size": 0,
+    "eos_token_id": 50257,
+    "exponential_decay_length_penalty": null,
+    "finetuning_task": null,
+    "forced_bos_token_id": null,
+    "forced_decoder_ids": [
+      [
+        1,
+        50259
+      ],
+      [
+        2,
+        50359
+      ],
+      [
+        3,
+        50363
+      ]
+    ],
+    "forced_eos_token_id": null,
+    "id2label": {
+      "0": "LABEL_0",
+      "1": "LABEL_1"
+    },
+    "init_std": 0.02,
+    "is_decoder": false,
+    "is_encoder_decoder": true,
+    "label2id": {
+      "LABEL_0": 0,
+      "LABEL_1": 1
+    },
+    "length_penalty": 1.0,
+    "mask_feature_length": 10,
+    "mask_feature_min_masks": 0,
+    "mask_feature_prob": 0.0,
+    "mask_time_length": 10,
+    "mask_time_min_masks": 2,
+    "mask_time_prob": 0.05,
+    "max_length": 448,
+    "max_source_positions": 1500,
+    "max_target_positions": 448,
+    "median_filter_width": 7,
+    "min_length": 0,
+    "model_type": "whisper",
+    "no_repeat_ngram_size": 0,
+    "num_beam_groups": 1,
+    "num_beams": 1,
+    "num_hidden_layers": 12,
+    "num_mel_bins": 80,
+    "num_return_sequences": 1,
+    "output_attentions": false,
+    "output_hidden_states": false,
+    "output_scores": false,
+    "pad_token_id": 50257,
+    "prefix": null,
+    "problem_type": null,
+    "pruned_heads": {},
+    "remove_invalid_values": false,
+    "repetition_penalty": 1.0,
+    "return_dict": true,
+    "return_dict_in_generate": false,
+    "scale_embedding": false,
+    "sep_token_id": null,
+    "suppress_tokens": [
+      1,
+      2,
+      7,
+      8,
+      9,
+      10,
+      14,
+      25,
+      26,
+      27,
+      28,
+      29,
+      31,
+      58,
+      59,
+      60,
+      61,
+      62,
+      63,
+      90,
+      91,
+      92,
+      93,
+      359,
+      503,
+      522,
+      542,
+      873,
+      893,
+      902,
+      918,
+      922,
+      931,
+      1350,
+      1853,
+      1982,
+      2460,
+      2627,
+      3246,
+      3253,
+      3268,
+      3536,
+      3846,
+      3961,
+      4183,
+      4667,
+      6585,
+      6647,
+      7273,
+      9061,
+      9383,
+      10428,
+      10929,
+      11938,
+      12033,
+      12331,
+      12562,
+      13793,
+      14157,
+      14635,
+      15265,
+      15618,
+      16553,
+      16604,
+      18362,
+      18956,
+      20075,
+      21675,
+      22520,
+      26130,
+      26161,
+      26435,
+      28279,
+      29464,
+      31650,
+      32302,
+      32470,
+      36865,
+      42863,
+      47425,
+      49870,
+      50254,
+      50258,
+      50360,
+      50361,
+      50362
+    ],
+    "task_specific_params": null,
+    "temperature": 1.0,
+    "tf_legacy_loss": false,
+    "tie_encoder_decoder": false,
+    "tie_word_embeddings": true,
+    "tokenizer_class": null,
+    "top_k": 50,
+    "top_p": 1.0,
+    "torch_dtype": "bfloat16",
+    "torchscript": false,
+    "typical_p": 1.0,
+    "use_bfloat16": false,
+    "use_cache": true,
+    "use_weighted_layer_sum": false,
+    "vocab_size": 51865
+  },
+  "audio_select_layer": -2,
+  "auto_map": {
+    "AutoConfig": "modeling_combine.MM_LLMs_Config",
+    "AutoModel": "modeling_combine.MM_LLMs"
+  },
+  "image_config": {
+    "_name_or_path": "google/siglip-base-patch16-384",
+    "add_cross_attention": false,
+    "architectures": [
+      "SiglipModel"
+    ],
+    "bad_words_ids": null,
+    "begin_suppress_tokens": null,
+    "bos_token_id": null,
+    "chunk_size_feed_forward": 0,
+    "cross_attention_hidden_size": null,
+    "decoder_start_token_id": null,
+    "diversity_penalty": 0.0,
+    "do_sample": false,
+    "early_stopping": false,
+    "encoder_no_repeat_ngram_size": 0,
+    "eos_token_id": null,
+    "exponential_decay_length_penalty": null,
+    "finetuning_task": null,
+    "forced_bos_token_id": null,
+    "forced_eos_token_id": null,
+    "id2label": {
+      "0": "LABEL_0",
+      "1": "LABEL_1"
+    },
+    "initializer_factor": 1.0,
+    "is_decoder": false,
+    "is_encoder_decoder": false,
+    "label2id": {
+      "LABEL_0": 0,
+      "LABEL_1": 1
+    },
+    "length_penalty": 1.0,
+    "max_length": 20,
+    "min_length": 0,
+    "model_type": "siglip",
+    "no_repeat_ngram_size": 0,
+    "num_beam_groups": 1,
+    "num_beams": 1,
+    "num_return_sequences": 1,
+    "output_attentions": false,
+    "output_hidden_states": false,
+    "output_scores": false,
+    "pad_token_id": null,
+    "prefix": null,
+    "problem_type": null,
+    "pruned_heads": {},
+    "remove_invalid_values": false,
+    "repetition_penalty": 1.0,
+    "return_dict": true,
+    "return_dict_in_generate": false,
+    "sep_token_id": null,
+    "suppress_tokens": null,
+    "task_specific_params": null,
+    "temperature": 1.0,
+    "text_config": {
+      "_name_or_path": "",
+      "add_cross_attention": false,
+      "architectures": null,
+      "attention_dropout": 0.0,
+      "bad_words_ids": null,
+      "begin_suppress_tokens": null,
+      "bos_token_id": 49406,
+      "chunk_size_feed_forward": 0,
+      "cross_attention_hidden_size": null,
+      "decoder_start_token_id": null,
+      "diversity_penalty": 0.0,
+      "do_sample": false,
+      "early_stopping": false,
+      "encoder_no_repeat_ngram_size": 0,
+      "eos_token_id": 49407,
+      "exponential_decay_length_penalty": null,
+      "finetuning_task": null,
+      "forced_bos_token_id": null,
+      "forced_eos_token_id": null,
+      "hidden_act": "gelu_pytorch_tanh",
+      "hidden_size": 768,
+      "id2label": {
+        "0": "LABEL_0",
+        "1": "LABEL_1"
+      },
+      "intermediate_size": 3072,
+      "is_decoder": false,
+      "is_encoder_decoder": false,
+      "label2id": {
+        "LABEL_0": 0,
+        "LABEL_1": 1
+      },
+      "layer_norm_eps": 1e-06,
+      "length_penalty": 1.0,
+      "max_length": 20,
+      "max_position_embeddings": 64,
+      "min_length": 0,
+      "model_type": "siglip_text_model",
+      "no_repeat_ngram_size": 0,
+      "num_attention_heads": 12,
+      "num_beam_groups": 1,
+      "num_beams": 1,
+      "num_hidden_layers": 12,
+      "num_return_sequences": 1,
+      "output_attentions": false,
+      "output_hidden_states": false,
+      "output_scores": false,
+      "pad_token_id": 1,
+      "prefix": null,
+      "problem_type": null,
+      "pruned_heads": {},
+      "remove_invalid_values": false,
+      "repetition_penalty": 1.0,
+      "return_dict": true,
+      "return_dict_in_generate": false,
+      "sep_token_id": null,
+      "suppress_tokens": null,
+      "task_specific_params": null,
+      "temperature": 1.0,
+      "tf_legacy_loss": false,
+      "tie_encoder_decoder": false,
+      "tie_word_embeddings": true,
+      "tokenizer_class": null,
+      "top_k": 50,
+      "top_p": 1.0,
+      "torch_dtype": null,
+      "torchscript": false,
+      "typical_p": 1.0,
+      "use_bfloat16": false,
+      "vocab_size": 32000
+    },
+    "tf_legacy_loss": false,
+    "tie_encoder_decoder": false,
+    "tie_word_embeddings": true,
+    "tokenizer_class": null,
+    "top_k": 50,
+    "top_p": 1.0,
+    "torch_dtype": "float32",
+    "torchscript": false,
+    "typical_p": 1.0,
+    "use_bfloat16": false,
+    "vision_config": {
+      "_name_or_path": "",
+      "add_cross_attention": false,
+      "architectures": null,
+      "attention_dropout": 0.0,
+      "bad_words_ids": null,
+      "begin_suppress_tokens": null,
+      "bos_token_id": null,
+      "chunk_size_feed_forward": 0,
+      "cross_attention_hidden_size": null,
+      "decoder_start_token_id": null,
+      "diversity_penalty": 0.0,
+      "do_sample": false,
+      "early_stopping": false,
+      "encoder_no_repeat_ngram_size": 0,
+      "eos_token_id": null,
+      "exponential_decay_length_penalty": null,
+      "finetuning_task": null,
+      "forced_bos_token_id": null,
+      "forced_eos_token_id": null,
+      "hidden_act": "gelu_pytorch_tanh",
+      "hidden_size": 768,
+      "id2label": {
+        "0": "LABEL_0",
+        "1": "LABEL_1"
+      },
+      "image_size": 384,
+      "intermediate_size": 3072,
+      "is_decoder": false,
+      "is_encoder_decoder": false,
+      "label2id": {
+        "LABEL_0": 0,
+        "LABEL_1": 1
+      },
+      "layer_norm_eps": 1e-06,
+      "length_penalty": 1.0,
+      "max_length": 20,
+      "min_length": 0,
+      "model_type": "siglip_vision_model",
+      "no_repeat_ngram_size": 0,
+      "num_attention_heads": 12,
+      "num_beam_groups": 1,
+      "num_beams": 1,
+      "num_channels": 3,
+      "num_hidden_layers": 12,
+      "num_return_sequences": 1,
+      "output_attentions": false,
+      "output_hidden_states": false,
+      "output_scores": false,
+      "pad_token_id": null,
+      "patch_size": 16,
+      "prefix": null,
+      "problem_type": null,
+      "pruned_heads": {},
+      "remove_invalid_values": false,
+      "repetition_penalty": 1.0,
+      "return_dict": true,
+      "return_dict_in_generate": false,
+      "sep_token_id": null,
+      "suppress_tokens": null,
+      "task_specific_params": null,
+      "temperature": 1.0,
+      "tf_legacy_loss": false,
+      "tie_encoder_decoder": false,
+      "tie_word_embeddings": true,
+      "tokenizer_class": null,
+      "top_k": 50,
+      "top_p": 1.0,
+      "torch_dtype": null,
+      "torchscript": false,
+      "typical_p": 1.0,
+      "use_bfloat16": false
+    }
+  },
+  "llm_config": {
+    "_name_or_path": "mesolitica/malaysian-mistral-7b-32k-instructions-v4",
+    "add_cross_attention": false,
+    "architectures": [
+      "MistralForCausalLM"
+    ],
+    "attention_dropout": 0.0,
+    "bad_words_ids": null,
+    "begin_suppress_tokens": null,
+    "bos_token_id": 1,
+    "chunk_size_feed_forward": 0,
+    "cross_attention_hidden_size": null,
+    "decoder_start_token_id": null,
+    "diversity_penalty": 0.0,
+    "do_sample": false,
+    "early_stopping": false,
+    "encoder_no_repeat_ngram_size": 0,
+    "eos_token_id": 2,
+    "exponential_decay_length_penalty": null,
+    "finetuning_task": null,
+    "forced_bos_token_id": null,
+    "forced_eos_token_id": null,
+    "hidden_act": "silu",
+    "hidden_size": 4096,
+    "id2label": {
+      "0": "LABEL_0",
+      "1": "LABEL_1"
+    },
+    "initializer_range": 0.02,
+    "intermediate_size": 14336,
+    "is_decoder": false,
+    "is_encoder_decoder": false,
+    "label2id": {
+      "LABEL_0": 0,
+      "LABEL_1": 1
+    },
+    "length_penalty": 1.0,
+    "max_length": 20,
+    "max_position_embeddings": 32768,
+    "min_length": 0,
+    "model_type": "mistral",
+    "no_repeat_ngram_size": 0,
+    "num_attention_heads": 32,
+    "num_beam_groups": 1,
+    "num_beams": 1,
+    "num_hidden_layers": 32,
+    "num_key_value_heads": 8,
+    "num_return_sequences": 1,
+    "output_attentions": false,
+    "output_hidden_states": false,
+    "output_scores": false,
+    "pad_token_id": null,
+    "prefix": null,
+    "problem_type": null,
+    "pruned_heads": {},
+    "remove_invalid_values": false,
+    "repetition_penalty": 1.0,
+    "return_dict": true,
+    "return_dict_in_generate": false,
+    "rms_norm_eps": 1e-05,
+    "rope_theta": 10000.0,
+    "sep_token_id": null,
+    "sliding_window": 4096,
+    "suppress_tokens": null,
+    "task_specific_params": null,
+    "temperature": 1.0,
+    "tf_legacy_loss": false,
+    "tie_encoder_decoder": false,
+    "tie_word_embeddings": false,
+    "tokenizer_class": null,
+    "top_k": 50,
+    "top_p": 1.0,
+    "torch_dtype": "bfloat16",
+    "torchscript": false,
+    "typical_p": 1.0,
+    "use_bfloat16": false,
+    "use_cache": true,
+    "vocab_size": 32004
+  },
+  "model_type": "mm_llms",
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.37.2",
+  "vision_select_layer": -2
+}

generation_config.json

@@ -0,0 +1,4 @@
+{
+  "_from_model_config": true,
+  "transformers_version": "4.37.2"
+}

model-00001-of-00004.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5161806048dd44f72c03b60315714b36fe294154deb1ccac549abb7f1e5bb5c6
+size 4961225036

model-00002-of-00004.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:22477ce687ca06bd01af0fe71049127e6ffb41331cec89d02bba68d288f8ca9d
+size 4999819752

model-00003-of-00004.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:93bda8dcfa1ec04e912d9378243dec58b8dc9985e04a60631bca7ba3d290c83c
+size 4915916584

model-00004-of-00004.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a435eaeed33c8d8f53e8104d759442ad6c36dab681955bd2ade80071f45096fe
+size 822175216

modeling_combine.py

@@ -0,0 +1,334 @@
+from collections import Counter, defaultdict
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch import Tensor
+from torch import nn
+from torch.nn import CrossEntropyLoss
+import copy
+import math
+from transformers.activations import gelu
+from typing import List, Optional, Tuple, Union
+from transformers.modeling_utils import PreTrainedModel, PretrainedConfig
+from transformers import CONFIG_MAPPING
+from transformers.modeling_outputs import BaseModelOutput
+from transformers import GenerationConfig
+from transformers import CLIPConfig, CLIPProcessor, CLIPModel, AutoModel
+from transformers import WhisperConfig, WhisperPreTrainedModel, WhisperModel
+from transformers import AutoConfig, AutoModelForCausalLM, LlamaConfig
+
+
+def most_frequent_element(tensor):
+    flattened_list = tensor.flatten().tolist()
+    counter = Counter(flattened_list)
+    most_common_element = counter.most_common(1)[0][1]
+
+    return most_common_element
+
+
+class MM_LLMs_Config(PretrainedConfig):
+    model_type = 'mm_llms'
+    is_composition = True
+
+    def __init__(
+        self,
+        image_config=None,
+        audio_config=None,
+        llm_config=None,
+        audio_select_layer=-2,
+        vision_select_layer=-2,
+        **kwargs
+    ):
+
+        self.image_config = image_config
+        self.audio_config = audio_config
+        self.llm_config = llm_config
+        self.audio_select_layer = audio_select_layer
+        self.vision_select_layer = vision_select_layer
+
+        if isinstance(self.image_config, dict):
+            image_config["model_type"] = (
+                image_config["model_type"] if "model_type" in image_config else "clip"
+            )
+            self.image_config = CONFIG_MAPPING[image_config["model_type"]](**image_config)
+        if isinstance(self.audio_config, dict):
+            audio_config["model_type"] = (
+                audio_config["model_type"] if "model_type" in audio_config else "whisper"
+            )
+            self.audio_config = CONFIG_MAPPING[audio_config["model_type"]](**audio_config)
+        if isinstance(self.llm_config, dict):
+            llm_config["model_type"] = llm_config["model_type"] if "model_type" in llm_config else "llama"
+            self.llm_config = CONFIG_MAPPING[llm_config["model_type"]](**llm_config)
+
+        super().__init__(**kwargs)
+
+
+class LlavaMultiModalProjector(nn.Module):
+    def __init__(self, in_hidden_size, out_hidden_size, conv_kernel=None, conv_stride=3):
+        super().__init__()
+
+        self.conv_kernel = conv_kernel
+
+        if conv_kernel:
+            self.linear_1 = nn.Conv1d(
+                in_hidden_size,
+                out_hidden_size,
+                kernel_size=conv_kernel,
+                stride=conv_stride)
+        else:
+            self.linear_1 = nn.Linear(
+                in_hidden_size,
+                out_hidden_size,
+                bias=True,
+            )
+        self.act = gelu
+        self.linear_2 = nn.Linear(
+            out_hidden_size,
+            out_hidden_size,
+            bias=True)
+
+    def forward(self, image_features):
+        hidden_states = self.linear_1(image_features)
+        if self.conv_kernel:
+            hidden_states = hidden_states.transpose(1, 2).contiguous()
+        hidden_states = self.act(hidden_states)
+        hidden_states = self.linear_2(hidden_states)
+        return hidden_states
+
+
+class MM_LLMs(PreTrainedModel):
+    config_class = MM_LLMs_Config
+    supports_gradient_checkpointing = True
+    _supports_flash_attn_2 = True
+
+    def __init__(self, config, flash_attention=False, dtype=torch.float32):
+        super().__init__(config)
+        self.config = config
+
+        self.image_encoder = AutoModel.from_config(config.image_config)
+
+        self.audio_encoder = AutoModel.from_config(
+            config.audio_config,
+            use_flash_attention_2=flash_attention,
+            torch_dtype=dtype,
+        )
+
+        self.llm = AutoModelForCausalLM.from_config(
+            config.llm_config,
+            use_flash_attention_2=flash_attention,
+            torch_dtype=dtype,
+        )
+
+        self.image_projector = LlavaMultiModalProjector(
+            config.image_config.vision_config.hidden_size,
+            config.llm_config.hidden_size
+        )
+        self.audio_projector = LlavaMultiModalProjector(
+            config.audio_config.d_model,
+            config.llm_config.hidden_size,
+            conv_kernel=40,
+            conv_stride=3,
+        )
+
+    def forward(self,
+                input_ids: torch.LongTensor = None,
+                image_index: torch.LongTensor = None,
+                audio_index: torch.LongTensor = None,
+                image_starts: torch.int = None,
+                image_ends: torch.int = None,
+                audio_starts: torch.int = None,
+                audio_ends: torch.int = None,
+                images: torch.FloatTensor = None,
+                audios: torch.FloatTensor = None,
+                attention_mask: Optional[torch.Tensor] = None,
+                position_ids: Optional[torch.LongTensor] = None,
+                past_key_values: Optional[List[torch.FloatTensor]] = None,
+                inputs_embeds: Optional[torch.FloatTensor] = None,
+                labels: Optional[torch.LongTensor] = None,
+                output_attentions: Optional[bool] = None,
+                output_hidden_states: Optional[bool] = None,
+                use_cache: Optional[bool] = None,
+                return_dict: Optional[bool] = None,
+                where_is_b=None,
+                where_is_k=None,
+                ls=None,
+                **kwargs):
+
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        images = images.type(self.image_encoder.dtype) if images is not None else None
+        audios = audios.type(self.audio_encoder.dtype) if audios is not None else None
+
+        print(
+            where_is_b,
+            where_is_k,
+            ls,
+            image_index,
+            audio_index,
+            input_ids.shape,
+            images.shape,
+            audios.shape,
+        )
+
+        model_inputs = self.prepare_inputs_for_generation(
+            input_ids=input_ids,
+            image_index=image_index,
+            audio_index=audio_index,
+            image_starts=image_starts,
+            image_ends=image_ends,
+            audio_starts=audio_starts,
+            audio_ends=audio_ends,
+            images=images,
+            audios=audios,
+            attention_mask=attention_mask,
+            labels=labels,
+            where_is_b=where_is_b,
+            where_is_k=where_is_k,
+            ls=ls)
+
+        outputs = self.llm(
+            inputs_embeds=model_inputs['inputs_embeds'],
+            attention_mask=model_inputs['attention_mask'],
+            labels=model_inputs['labels'],
+            return_dict=return_dict)
+
+        return outputs
+
+    def prepare_inputs_for_generation(
+            self,
+            input_ids,
+            past_key_values=None,
+            inputs_embeds=None,
+            images=None,
+            audios=None,
+            audio_starts=None,
+            audio_ends=None,
+            image_starts=None,
+            image_ends=None,
+            attention_mask=None,
+            labels=None,
+            audio_index=None,
+            image_index=None,
+            where_is_b=None,
+            where_is_k=None,
+            ls=None,
+            inference=False,
+            **kwargs):
+
+        image_features = self.encode_image(
+            images) if images is not None else None
+        audio_features = self.encode_audio(
+            audios) if audios is not None else None
+        embed_tokens = self.llm.model.embed_tokens
+        text_embeddings = embed_tokens(input_ids)
+        batch_size = text_embeddings.shape[0]
+        seq_len = text_embeddings.shape[1]
+        embed_dim = text_embeddings.shape[2]
+
+        if len(audio_index):
+            max_count_audio = most_frequent_element(audio_index)
+        else:
+            max_count_audio = 0
+        if len(image_index):
+            max_count_image = most_frequent_element(image_index)
+        else:
+            max_count_image = 0
+
+        if audio_features is not None:
+            seq_audio = audio_features.shape[1]
+        else:
+            seq_audio = 0
+
+        if image_features is not None:
+            seq_image = image_features.shape[1]
+        else:
+            seq_image = 0
+
+        audio_len = seq_audio * max_count_audio
+        image_len = seq_image * max_count_image
+
+        new_len = text_embeddings.shape[1] + audio_len + image_len
+        final_embedding = torch.zeros(
+            batch_size, new_len, embed_dim,
+            device=text_embeddings.device,
+            dtype=text_embeddings.dtype
+        )
+        final_embedding[:, :seq_len] = text_embeddings
+        final_attention_mask = torch.zeros(
+            batch_size, new_len,
+            device=attention_mask.device,
+            dtype=attention_mask.dtype
+        )
+        final_attention_mask[:, :seq_len] = attention_mask
+        if labels is not None:
+            final_labels = torch.full(
+                (batch_size, new_len),
+                -100,
+                device=labels.device,
+                dtype=labels.dtype
+            )
+            final_labels[:, :seq_len] = labels
+        else:
+            final_labels = None
+
+        image_id = int(image_starts[0])
+        audio_id = int(audio_starts[0])
+
+        positions = defaultdict(int)
+        b_image = 0
+        b_audio = 0
+
+        for i in range(len(where_is_b)):
+            b, k = where_is_b[i], where_is_k[i]
+            int_b = int(b)
+            int_k = int(k)
+            l = int(ls[i])
+            if l == image_id:
+                f = image_features[b_image]
+                b_image += 1
+            if l == audio_id:
+                f = audio_features[b_audio]
+                b_audio += 1
+
+            c = torch.cat([final_embedding[b, :int_k + 1 + positions[int_b]],
+                          f, text_embeddings[b, k + 1:]])
+            final_embedding[b, :len(c)] = c
+            final_attention_mask[b, :len(c)] = 1.0
+
+            if labels is not None:
+                ignore = torch.tensor([-100] * len(f), device=labels.device)
+                c_label = torch.cat(
+                    [final_labels[b, :int_k + 1 + positions[int_b]], ignore, labels[b, k + 1:]])
+                final_labels[b, :len(c)] = c_label
+
+            positions[int_b] += len(f)
+
+        if not inference:
+            final_attention_mask[:, :seq_audio + seq_image + 2] = 0.0
+
+        if labels is not None:
+            final_labels[:, :seq_audio + seq_image + 2] = -100
+
+        model_inputs = {
+            "input_ids": input_ids,
+            "inputs_embeds": final_embedding,
+            "use_cache": kwargs.get("use_cache"),
+            "attention_mask": final_attention_mask,
+            "labels": final_labels,
+        }
+        return model_inputs
+
+    def encode_audio(self, audios):
+        encoded = self.audio_encoder.encoder(audios, output_hidden_states=True)
+        encoded = encoded.hidden_states[self.config.audio_select_layer]
+        audio_features = self.audio_projector(encoded.transpose(1, 2).contiguous())
+        return audio_features
+
+    def encode_image(self, images):
+        if self.config.vision_select_layer is not None:
+            encoded = self.image_encoder.vision_model(images, output_hidden_states=True)
+            encoded = encoded.hidden_states[self.config.vision_select_layer]
+        else:
+            encoded = self.image_encoder.vision_model(images)[0]
+        image_features = self.image_projector(encoded)
+        return image_features