Make More Generic; Reduce Config Size

config.json

@@ -1,138 +1,12 @@
 {
   "model_type": "diva",
+  "reference_encoder": "openai/whisper-large-v3",
+  "reference_decoder": "meta-llama/Meta-Llama-3-8B-Instruct",
   "architectures": [ "DiVAModel" ],
   "auto_map": {
     "AutoConfig": "configuring_diva.DiVAConfig",
     "AutoModel": "modeling_diva.DiVAModel"
   },
   "vocab_size": 128256,
-  "decoder": {
-    "architectures": [
-      "LlamaForCausalLM"
-    ],
-    "attention_bias": false,
-    "attention_dropout": 0,
-    "bos_token_id": 128000,
-    "eos_token_id": 128001,
-    "hidden_act": "silu",
-    "hidden_size": 4096,
-    "initializer_range": 0.02,
-    "intermediate_size": 14336,
-    "max_position_embeddings": 8192,
-    "model_type": "llama",
-    "num_attention_heads": 32,
-    "num_hidden_layers": 32,
-    "num_key_value_heads": 8,
-    "pretraining_tp": 1,
-    "rms_norm_eps": 1e-05,
-    "rope_scaling": null,
-    "rope_theta": 500000,
-    "tie_word_embeddings": false,
-    "torch_dtype": "bfloat16",
-    "transformers_version": "4.40.0.dev0",
-    "use_cache": true,
-    "vocab_size": 128256
-  },
-  "encoder": {
-    "_name_or_path": "openai/whisper-large-v3",
-    "activation_dropout": 0,
-    "activation_function": "gelu",
-    "add_cross_attention": false,
-    "apply_spec_augment": false,
-    "architectures": [
-      "WhisperForConditionalGeneration"
-    ],
-    "attention_dropout": 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": 1280,
-    "decoder_attention_heads": 20,
-    "decoder_ffn_dim": 5120,
-    "decoder_layerdrop": 0,
-    "decoder_layers": 32,
-    "decoder_start_token_id": 50258,
-    "diversity_penalty": 0,
-    "do_sample": false,
-    "dropout": 0,
-    "early_stopping": false,
-    "encoder_attention_heads": 20,
-    "encoder_ffn_dim": 5120,
-    "encoder_layerdrop": 0,
-    "encoder_layers": 32,
-    "encoder_no_repeat_ngram_size": 0,
-    "eos_token_id": 50257,
-    "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"
-    },
-    "init_std": 0.02,
-    "is_decoder": false,
-    "is_encoder_decoder": true,
-    "label2id": {
-      "LABEL_0": 0,
-      "LABEL_1": 1
-    },
-    "length_penalty": 1,
-    "mask_feature_length": 10,
-    "mask_feature_min_masks": 0,
-    "mask_feature_prob": 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": 32,
-    "num_mel_bins": 128,
-    "num_return_sequences": 1,
-    "output_attentions": false,
-    "output_hidden_states": false,
-    "output_scores": false,
-    "pad_token_id": 50256,
-    "prefix": null,
-    "problem_type": null,
-    "pruned_heads": {},
-    "remove_invalid_values": false,
-    "repetition_penalty": 1,
-    "return_dict": true,
-    "return_dict_in_generate": false,
-    "scale_embedding": false,
-    "sep_token_id": null,
-    "suppress_tokens": null,
-    "task_specific_params": null,
-    "temperature": 1,
-    "tf_legacy_loss": false,
-    "tie_encoder_decoder": false,
-    "tie_word_embeddings": true,
-    "tokenizer_class": null,
-    "top_k": 50,
-    "top_p": 1,
-    "torch_dtype": "float16",
-    "torchscript": false,
-    "transformers_version": "4.38.2",
-    "typical_p": 1,
-    "use_bfloat16": false,
-    "use_cache": true,
-    "use_weighted_layer_sum": false,
-    "vocab_size": 51866
-  },
-  "time_dialation": 4,
   "transformers_version": "4.38.2"
 }

modeling_diva.py

@@ -10,13 +10,13 @@ import torch.nn.functional as F
 from datasets import Audio
 from safetensors.torch import load, load_model
 from torch import nn
-from .configuring_diva import DiVAConfig
+from configuring_diva import DiVAConfig
 from transformers import (
     AutoProcessor,
     AutoTokenizer,
-    LlamaForCausalLM,
+    AutoModelForCausalLM,
     PreTrainedModel,
-    WhisperForConditionalGeneration,
+    WhisperModel,
 )
 
 
@@ -51,11 +51,9 @@ class DiVAModel(PreTrainedModel):
         super().__init__(DiVAConfig.from_dict(config_dict))
         if speech_encoder_device is None:
             speech_encoder_device = "cuda:0"
-        whisper = WhisperForConditionalGeneration.from_pretrained(
-            "openai/whisper-large-v3"
-        )
+        whisper = WhisperModel.from_pretrained(config_dict["reference_encoder"])
         connector = WhisperConnector()
-        connector.decoder = copy.deepcopy(whisper.model.decoder)
+        connector.decoder = copy.deepcopy(whisper.decoder)
         if via_path is not None:
             with open(via_path, "rb") as f:
                 sd = load(f.read())
@@ -83,25 +81,25 @@ class DiVAModel(PreTrainedModel):
             )
 
         self.connector = connector.to(speech_encoder_device)
-        self.whisper_encoder = whisper.model.encoder.to(speech_encoder_device)
-        self.llama_decoder = LlamaForCausalLM.from_pretrained(
-            "meta-llama/Meta-Llama-3-8B-Instruct",
+        self.whisper_encoder = whisper.encoder.to(speech_encoder_device)
+        self.llm_decoder = AutoModelForCausalLM.from_pretrained(
+            config_dict["reference_decoder"],
             device_map=device_map,
             torch_dtype=torch.float16,
         )
-        self.processor = AutoProcessor.from_pretrained("openai/whisper-large-v3")
+        self.processor = AutoProcessor.from_pretrained(config_dict["reference_encoder"])
         self.tokenizer = AutoTokenizer.from_pretrained("WillHeld/via-llama")
         self.prefix = torch.tensor([128000, 128006, 882, 128007, 271]).to(
-            self.llama_decoder.model.embed_tokens.weight.device
+            self.llm_decoder.model.embed_tokens.weight.device
         )
 
         self.pre_user_suffix = torch.tensor(
             self.tokenizer.encode(
                 "<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n"
             )
-        ).to(self.llama_decoder.model.embed_tokens.weight.device)
+        ).to(self.llm_decoder.model.embed_tokens.weight.device)
         self.final_header = torch.tensor([128009, 128006, 78191, 128007, 271]).to(
-            self.llama_decoder.model.embed_tokens.weight.device
+            self.llm_decoder.model.embed_tokens.weight.device
         )
         self.speech_encoder_device = speech_encoder_device
 
@@ -161,18 +159,18 @@ class DiVAModel(PreTrainedModel):
         ]
         virt_tokens = self.connector(
             hidden_states,
-            output_device=self.llama_decoder.model.embed_tokens.weight.device,
+            output_device=self.llm_decoder.model.embed_tokens.weight.device,
         ).squeeze()
 
-        prefix_embed = self.llama_decoder.model.embed_tokens(prefix_text_tokens)
-        suffix_embed = self.llama_decoder.model.embed_tokens(suffix_text_tokens)
+        prefix_embed = self.llm_decoder.model.embed_tokens(prefix_text_tokens)
+        suffix_embed = self.llm_decoder.model.embed_tokens(suffix_text_tokens)
         inputs_embeds = torch.cat(
             [prefix_embed, virt_tokens, suffix_embed], axis=0
         ).unsqueeze(0)
 
-        outputs = self.llama_decoder(
+        outputs = self.llm_decoder(
             inputs_embeds=inputs_embeds.to(
-                self.llama_decoder.model.embed_tokens.weight.device
+                self.llm_decoder.model.embed_tokens.weight.device
             ).half(),
             return_dict=True,
             output_hidden_states=True,
@@ -197,7 +195,7 @@ class DiVAModel(PreTrainedModel):
         ]
         virt_tokens = self.connector(
             hidden_states,
-            output_device=self.llama_decoder.model.embed_tokens.weight.device,
+            output_device=self.llm_decoder.model.embed_tokens.weight.device,
         )
         bsz = virt_tokens.shape[0]
 
@@ -227,9 +225,9 @@ class DiVAModel(PreTrainedModel):
             )
         else:
             prefix = self.prefix
-        prefix_embed = self.llama_decoder.model.embed_tokens(prefix).expand(bsz, -1, -1)
+        prefix_embed = self.llm_decoder.model.embed_tokens(prefix).expand(bsz, -1, -1)
         suffix = self.final_header
-        suffix_embed = self.llama_decoder.model.embed_tokens(suffix).expand(bsz, -1, -1)
+        suffix_embed = self.llm_decoder.model.embed_tokens(suffix).expand(bsz, -1, -1)
         inputs_embeds = torch.cat([prefix_embed, virt_tokens, suffix_embed], axis=1)
         outs = [[] for i in range(bsz)]
         complete = [False] * bsz
@@ -238,9 +236,9 @@ class DiVAModel(PreTrainedModel):
         i = 0
         while not all(complete) and len(outs[0]) < max_new_tokens:
             past_key_values = outputs.past_key_values if outputs else None
-            outputs = self.llama_decoder(
+            outputs = self.llm_decoder(
                 inputs_embeds=inputs_embeds.to(
-                    self.llama_decoder.model.embed_tokens.weight.device
+                    self.llm_decoder.model.embed_tokens.weight.device
                 ).half(),
                 return_dict=True,
                 output_hidden_states=True,
@@ -268,7 +266,7 @@ class DiVAModel(PreTrainedModel):
                 if out == 128009:
                     complete[token_index] = True
 
-            next_embed = self.llama_decoder.model.embed_tokens(greedy.reshape(-1, 1))
+            next_embed = self.llm_decoder.model.embed_tokens(greedy.reshape(-1, 1))
             inputs_embeds = next_embed
         return self.tokenizer.batch_decode(outs, skip_special_tokens=True)
 
@@ -287,7 +285,7 @@ class DiVAModel(PreTrainedModel):
         ]
         virt_tokens = self.connector(
             hidden_states,
-            output_device=self.llama_decoder.model.embed_tokens.weight.device,
+            output_device=self.llm_decoder.model.embed_tokens.weight.device,
         ).squeeze()
 
         if text_prompt != None and text_prompt != "":
@@ -300,9 +298,9 @@ class DiVAModel(PreTrainedModel):
             )
         else:
             prefix = self.prefix
-        prefix_embed = self.llama_decoder.model.embed_tokens(prefix)
+        prefix_embed = self.llm_decoder.model.embed_tokens(prefix)
         suffix = self.final_header
-        suffix_embed = self.llama_decoder.model.embed_tokens(suffix)
+        suffix_embed = self.llm_decoder.model.embed_tokens(suffix)
         inputs_embeds = torch.cat(
             [prefix_embed, virt_tokens, suffix_embed], axis=0
         ).unsqueeze(0)
@@ -312,9 +310,9 @@ class DiVAModel(PreTrainedModel):
         i = 0
         while greedy != 128009 and len(outs) < max_new_tokens:
             past_key_values = outputs.past_key_values if outputs else None
-            outputs = self.llama_decoder(
+            outputs = self.llm_decoder(
                 inputs_embeds=inputs_embeds.to(
-                    self.llama_decoder.model.embed_tokens.weight.device
+                    self.llm_decoder.model.embed_tokens.weight.device
                 ).half(),
                 return_dict=True,
                 output_hidden_states=True,
@@ -337,7 +335,7 @@ class DiVAModel(PreTrainedModel):
             else:
                 greedy = next_token_logits.argmax()
             outs.append(greedy)
-            next_embed = self.llama_decoder.model.embed_tokens(greedy.reshape(1, 1))
+            next_embed = self.llm_decoder.model.embed_tokens(greedy.reshape(1, 1))
             inputs_embeds = next_embed
             yield self.tokenizer.decode(outs, skip_special_tokens=True).replace(
                 "<|eot_id|>", ""