Update modeling_typhoonaudio.py

refactor generate() and streaming for talk areana

modeling_typhoonaudio.py

@@ -15,10 +15,9 @@ from transformers import (
     WhisperModel,
     PreTrainedModel,
     AutoTokenizer, 
-    AutoModelForCausalLM
+    AutoModelForCausalLM,
+    TextIteratorStreamer
 )
-import soundfile as sf
-import librosa
 from .configuration_typhoonaudio import TyphoonAudioConfig
 # ---------------------------------------------------- # 
 # QFormer: https://github.com/huggingface/transformers/blob/v4.15.0/src/transformers/models/bert
@@ -156,35 +155,26 @@ class TyphoonAudio(PreTrainedModel):
 
     def generate(
         self,
-        wav_path,
+        audio,
         prompt,
         prompt_pattern,
-        device='cuda:0',
-        max_length=150,
+        max_new_tokens=1024,
         num_beams=4,
         do_sample=True,
-        min_length=1,
         top_p=0.9,
         repetition_penalty=1.0,
         length_penalty=1.0,
         temperature=1.0,
         streamer=None
-    ):
-        # read wav
-        wav, sr = sf.read(wav_path)
-        if len(wav.shape) == 2:
-            wav = wav[:, 0]
-        if len(wav) > 30 * sr:
-            wav = wav[: 30 * sr]
-        if sr != 16000:
-            wav = librosa.resample(wav, orig_sr=sr, target_sr=16000, res_type="fft")
-        
+    ):      
+        device = self.llama_model.device
+
         # whisper
-        spectrogram = self.feature_extractor(wav, return_tensors="pt", sampling_rate=16000).input_features.to(device).to(self.torch_dtype) # [1, 80, 3000]
+        spectrogram = self.feature_extractor(audio, return_tensors="pt", sampling_rate=16000).input_features.to(device).to(self.torch_dtype) # [1, 80, 3000]
         speech_embeds = self.speech_encoder(spectrogram, return_dict=True).last_hidden_state
 
         # beats
-        raw_wav = torch.from_numpy(wav).to(device).unsqueeze(0)
+        raw_wav = torch.from_numpy(audio).to(device).unsqueeze(0)
         audio_padding_mask = torch.zeros(raw_wav.shape, device=device).bool()
         audio_embeds, _ = self.beats.extract_features(raw_wav, padding_mask=audio_padding_mask, feature_only=True, torch_dtype=self.torch_dtype)
 
@@ -256,10 +246,9 @@ class TyphoonAudio(PreTrainedModel):
         # generate
         output = self.llama_model.generate(
             inputs_embeds=embeds,
-            max_length=max_length,
+            max_new_tokens=max_new_tokens,
             num_beams=num_beams,
             do_sample=do_sample,
-            min_length=min_length,
             top_p=top_p,
             repetition_penalty=repetition_penalty,
             length_penalty=length_penalty,
@@ -273,6 +262,38 @@ class TyphoonAudio(PreTrainedModel):
         output_text = self.llama_tokenizer.batch_decode(output, add_special_tokens=False, skip_special_tokens=True)
         return output_text[0]
 
+    def generate_stream(
+        self,
+        audio,
+        prompt,
+        prompt_pattern="<|start_header_id|>user<|end_header_id|>\n\n<Speech><SpeechHere></Speech> {}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n",
+        max_new_tokens=1024,
+        do_sample=True,
+        top_p=0.9,
+        repetition_penalty=1.0,
+        length_penalty=1.0,
+        temperature=1.0,
+    ):
+        streamer = TextIteratorStreamer(self.llama_tokenizer)
+        _ = self.generate(
+            audio=audio,
+            prompt=prompt,
+            prompt_pattern=prompt_pattern,
+            do_sample=do_sample,
+            max_new_tokens=max_new_tokens,
+            temperature=temperature,
+            top_p=top_p,
+            repetition_penalty=repetition_penalty,
+            length_penalty=length_penalty,
+            streamer=streamer,
+            num_beams=1,
+        )
+        response = ""
+        for new_tokens in streamer:
+            response += new_tokens.replace("<|eot_id|>", "").replace("<|end_of_text|>", "")
+            yield response
+        return response
+
     def init_speech_Qformer(self, num_query_token, speech_width, num_hidden_layers=2, torch_dtype="float16"):
         encoder_config = BertConfig()
         encoder_config.num_hidden_layers = num_hidden_layers