Add Batch Support

modeling_diva.py

@@ -44,7 +44,7 @@ class WhisperConnector(nn.Module):
 
 class DiVAModel(PreTrainedModel):
     config_class = DiVAConfig
-    
+
     def __init__(
         self, via_path=None, config_dict={}, device_map=None, speech_encoder_device=None
     ):
@@ -105,10 +105,9 @@ class DiVAModel(PreTrainedModel):
         )
         self.speech_encoder_device = speech_encoder_device
 
-
-    def can_generate(cls): 
+    def can_generate(cls):
         return False
-    
+
     @classmethod
     def from_pretrained(
         cls,
@@ -182,8 +181,14 @@ class DiVAModel(PreTrainedModel):
 
         return outputs
 
+    @torch.no_grad()
     def generate(
-        self, audio, text_prompt, do_sample=False, logits_processor=None, max_new_tokens=128
+        self,
+        audio,
+        text_prompt=None,
+        do_sample=False,
+        logits_processor=None,
+        max_new_tokens=128,
     ):
         inputs = self.processor(audio, return_tensors="pt", sampling_rate=16_000)
         input_features = inputs.input_features.to(self.speech_encoder_device)
@@ -193,29 +198,45 @@ class DiVAModel(PreTrainedModel):
         virt_tokens = self.connector(
             hidden_states,
             output_device=self.llama_decoder.model.embed_tokens.weight.device,
-        ).squeeze()
+        )
+        bsz = virt_tokens.shape[0]
 
         if text_prompt != None and text_prompt != "":
             user_prompt_text = torch.tensor(
-                self.tokenizer(text_prompt, add_special_tokens=False)["input_ids"],
+                self.tokenizer(
+                    text_prompt,
+                    add_special_tokens=False,
+                    padding=True,
+                    padding_side="right",
+                )["input_ids"],
                 device=self.pre_user_suffix.device,
             )
             prefix = torch.cat(
-                [self.pre_user_suffix, user_prompt_text, self.prefix], axis=0
+                [
+                    self.pre_user_suffix.expand(
+                        bsz,
+                        -1,
+                    ),
+                    user_prompt_text,
+                    self.prefix.expand(
+                        bsz,
+                        -1,
+                    ),
+                ],
+                axis=1,
             )
         else:
             prefix = self.prefix
-        prefix_embed = self.llama_decoder.model.embed_tokens(prefix)
+        prefix_embed = self.llama_decoder.model.embed_tokens(prefix).expand(bsz, -1, -1)
         suffix = self.final_header
-        suffix_embed = self.llama_decoder.model.embed_tokens(suffix)
-        inputs_embeds = torch.cat(
-            [prefix_embed, virt_tokens, suffix_embed], axis=0
-        ).unsqueeze(0)
-        outs = []
+        suffix_embed = self.llama_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
         outputs = None
         greedy = 1
         i = 0
-        while greedy != 128009 and len(outs) < max_new_tokens:
+        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(
                 inputs_embeds=inputs_embeds.to(
@@ -225,7 +246,7 @@ class DiVAModel(PreTrainedModel):
                 output_hidden_states=True,
                 past_key_values=past_key_values,
             )
-            next_token_logits = outputs.logits[-1, -1, :]
+            next_token_logits = outputs.logits[:, -1, :]
 
             if logits_processor:
                 local_outs = torch.tensor(outs) if outs != [] else suffix
@@ -240,16 +261,23 @@ class DiVAModel(PreTrainedModel):
                 probs = F.softmax(logits, dim=-1)
                 greedy = torch.multinomial(probs, num_samples=1)[0]
             else:
-                greedy = next_token_logits.argmax()
-            outs.append(greedy)
-            next_embed = self.llama_decoder.model.embed_tokens(greedy.reshape(1, 1))
+                greedy = next_token_logits.argmax(dim=-1)
+            for token_index, out in enumerate(greedy.flatten().tolist()):
+                outs[token_index].append(out)
+                if out == 128009:
+                    complete[token_index] = True
+
+            next_embed = self.llama_decoder.model.embed_tokens(greedy.reshape(-1, 1))
             inputs_embeds = next_embed
-        return self.tokenizer.decode(outs, skip_special_tokens=True).replace(
-            "<|eot_id|>", ""
-        )
+        return self.tokenizer.batch_decode(outs, skip_special_tokens=True)
 
     def generate_stream(
-        self, audio, text_prompt, do_sample=False, logits_processor=None, max_new_tokens=128
+        self,
+        audio,
+        text_prompt,
+        do_sample=False,
+        logits_processor=None,
+        max_new_tokens=128,
     ):
         inputs = self.processor(audio, return_tensors="pt", sampling_rate=16_000)
         input_features = inputs.input_features.to(self.whisper_encoder.device)
@@ -284,7 +312,7 @@ class DiVAModel(PreTrainedModel):
         while greedy != 128009 and len(outs) < max_new_tokens:
             past_key_values = outputs.past_key_values if outputs else None
             outputs = self.llama_decoder(
-                    inputs_embeds=inputs_embeds.to(
+                inputs_embeds=inputs_embeds.to(
                     self.llama_decoder.model.embed_tokens.weight.device
                 ).half(),
                 return_dict=True,
@@ -310,5 +338,9 @@ class DiVAModel(PreTrainedModel):
             outs.append(greedy)
             next_embed = self.llama_decoder.model.embed_tokens(greedy.reshape(1, 1))
             inputs_embeds = next_embed
-            yield self.tokenizer.decode(outs, skip_special_tokens=True).replace("<|eot_id|>", "")
-        return self.tokenizer.decode(outs, skip_special_tokens=True).replace("<|eot_id|>", "")
+            yield self.tokenizer.decode(outs, skip_special_tokens=True).replace(
+                "<|eot_id|>", ""
+            )
+        return self.tokenizer.decode(outs, skip_special_tokens=True).replace(
+            "<|eot_id|>", ""
+        )

test.py

@@ -0,0 +1,28 @@
+from transformers import AutoModel
+import librosa
+import wget
+from modeling_diva import DiVAModel
+
+filename = wget.download(
+    "https://github.com/ffaisal93/SD-QA/raw/refs/heads/master/dev/eng/irl/wav_eng/-1008642825401516622.wav"
+)
+
+speech_data, _ = librosa.load(filename, sr=16_000)
+
+model = DiVAModel.from_pretrained("./")
+
+print(model.generate([speech_data]))
+print(model.generate([speech_data], ["Reply Briefly Like A Pirate"]))
+
+filename = wget.download(
+    "https://github.com/ffaisal93/SD-QA/raw/refs/heads/master/dev/eng/irl/wav_eng/-2426554427049983479.wav"
+)
+
+speech_data2, _ = librosa.load(filename, sr=16_000)
+
+print(
+    model.generate(
+        [speech_data, speech_data2],
+        ["Reply Briefly Like A Pirate", "Reply Briefly Like A New Yorker"],
+    )
+)