generate_meta update

config.json

@@ -1,13 +1,18 @@
 {
+  "architectures": [
+    "SpeechLLMModel"
+  ],
   "audio_enc_dim": 1024,
   "audio_encoder_name": "microsoft/wavlm-large",
   "audio_processor_name": "microsoft/wavlm-base",
   "auto_map": {
-    "AutoConfig": "config.SpeechLLMModelConfig"
+    "AutoConfig": "config.SpeechLLMModelConfig",
+    "AutoModel": "model.SpeechLLMModel"
   },
   "llm_dim": 2048,
   "llm_model_checkpoint": "hf_repo/llm_model_checkpoint",
   "llm_model_name": "TinyLlama/TinyLlama-1.1B-Chat-v1.0",
   "model_type": "custom_model",
+  "torch_dtype": "float32",
   "transformers_version": "4.41.2"
 }

model.py

@@ -61,11 +61,11 @@ class SpeechLLMModel(PreTrainedModel):
         self.llm_model = get_peft_model(self.llm_model, peft_config)
         self.llm_model = self.llm_model.merge_and_unload()
 
-    def encode(self, mel, pre_tokenized_ids, post_tokenized_ids, output_tokenized_ids):
-        batch_size = mel.shape[0]
+    def encode(self, speech, pre_tokenized_ids, post_tokenized_ids, output_tokenized_ids):
+        batch_size = speech.shape[0]
 
         with torch.no_grad():
-            speech_embeds = self.audio_encoder(mel)
+            speech_embeds = self.audio_encoder(speech)
             speech_embeds = self.connector(speech_embeds)
             
         embedder = self.llm_model.model.embed_tokens
@@ -83,12 +83,12 @@ class SpeechLLMModel(PreTrainedModel):
         ], 1).to(combined_embeds.device).to(torch.int64)
         return combined_embeds, atts, label_ids
 
-    def forward(self, wav_tensor, pre_tokenized_ids, post_tokenized_ids, output_tokenized_ids, attention_mask=None):
-        combined_embeds, atts, label_ids = self.encode(wav_tensor, pre_tokenized_ids, post_tokenized_ids, output_tokenized_ids)
+    def forward(self, audio_tensor, pre_tokenized_ids, post_tokenized_ids, output_tokenized_ids, attention_mask=None):
+        combined_embeds, atts, label_ids = self.encode(audio_tensor, pre_tokenized_ids, post_tokenized_ids, output_tokenized_ids)
         outputs = self.llm_model(inputs_embeds=combined_embeds, attention_mask=attention_mask)
         return outputs
     
-    def generate_meta(self, audio_path, instruction="Give me the following information about the audio [Transcript]", max_new_tokens=2000):
+    def generate_meta(self, audio_path=None, audio_tensor=None, instruction="Give me the following information about the audio [Transcript]", max_new_tokens=2000):
         device = self.audio_encoder.return_device()
         pre_speech_prompt = f'''Instruction:
 {instruction}
@@ -101,14 +101,15 @@ Output:'''
         output_prompt = '\n<s>'
 
         with torch.no_grad():
-            wav_tensor, sr = torchaudio.load(audio_path)
-            wav_tensor = self.audio_processor(wav_tensor.squeeze(), return_tensors="pt", sampling_rate=16000).input_values
+            if audio_tensor == None and audio_path != None:
+                audio_tensor, sr = torchaudio.load(audio_path)
+            audio_tensor = self.audio_processor(audio_tensor.squeeze(), return_tensors="pt", sampling_rate=16000).input_values
 
             pre_tokenized_ids = self.llm_tokenizer(pre_speech_prompt, padding="do_not_pad", return_tensors='pt', truncation=False, add_special_tokens=False)["input_ids"]
             post_tokenized_ids = self.llm_tokenizer(post_speech_prompt, padding="do_not_pad", return_tensors='pt', truncation=False, add_special_tokens=False)["input_ids"]
             output_tokenized_ids = self.llm_tokenizer(output_prompt, padding="do_not_pad", return_tensors='pt', truncation=False, add_special_tokens=False)["input_ids"]
 
-            combined_embeds, atts, label_ids = self.encode(wav_tensor.to(device), pre_tokenized_ids.to(device), post_tokenized_ids.to(device), output_tokenized_ids.to(device))
+            combined_embeds, atts, label_ids = self.encode(audio_tensor.to(device), pre_tokenized_ids.to(device), post_tokenized_ids.to(device), output_tokenized_ids.to(device))
 
             out = self.llm_model.generate(
                 inputs_embeds=combined_embeds,