Update modeling_diva.py

modeling_diva.py

@@ -277,6 +277,8 @@ class DiVAModel(PreTrainedModel):
         do_sample=False,
         logits_processor=None,
         max_new_tokens=128,
+        return_outputs=False,
+        init_outputs=None,
     ):
         inputs = self.processor(audio, return_tensors="pt", sampling_rate=16_000)
         input_features = inputs.input_features.to(self.whisper_encoder.device)
@@ -305,7 +307,7 @@ class DiVAModel(PreTrainedModel):
             [prefix_embed, virt_tokens, suffix_embed], axis=0
         ).unsqueeze(0)
         outs = []
-        outputs = None
+        outputs = init_outputs
         greedy = 1
         i = 0
         while greedy != 128009 and len(outs) < max_new_tokens:
@@ -337,9 +339,19 @@ class DiVAModel(PreTrainedModel):
             outs.append(greedy)
             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(
+            if not return_outputs:
+                yield self.tokenizer.decode(outs, skip_special_tokens=True).replace(
+                    "<|eot_id|>", ""
+                )
+            else:
+                yield (self.tokenizer.decode(outs, skip_special_tokens=True).replace(
+                    "<|eot_id|>", ""
+                ), outputs)
+        if not return_outputs:
+            return self.tokenizer.decode(outs, skip_special_tokens=True).replace(
                 "<|eot_id|>", ""
             )
-        return self.tokenizer.decode(outs, skip_special_tokens=True).replace(
-            "<|eot_id|>", ""
-        )
+        else:
+            return (self.tokenizer.decode(outs, skip_special_tokens=True).replace(
+                "<|eot_id|>", ""
+            ), outputs)