adding new version of whisper

app.py

@@ -1,18 +1,18 @@
 
 import gradio as gr
-from whisper2 import generate
+from whisper import generate
 from AinaTheme import theme
 
 MODEL_NAME = "/whisper-large-v3"
+USE_V4 = False
 
 
-
-def transcribe(inputs):
+def transcribe(inputs, use_v4):
     if inputs is None:
         raise gr.Error("Cap fitxer d'àudio introduit! Si us plau pengeu un fitxer "\
                        "o enregistreu un àudio abans d'enviar la vostra sol·licitud")
 
-    return generate(audio=inputs)
+    return generate(audio_path=inputs, use_v4=use_v4)
 
 
 description_string = "Transcripció automàtica de micròfon o de fitxers d'àudio.\n Aquest demostrador s'ha desenvolupat per"\
@@ -22,7 +22,8 @@ description_string = "Transcripció automàtica de micròfon o de fitxers d'àud
 
 def clear():
      return (
-          None
+          None,
+          USE_V4
      )
 
 
@@ -30,6 +31,7 @@ with gr.Blocks(theme=theme) as demo:
     gr.Markdown(description_string)
     with gr.Row():
         with gr.Column(scale=1):
+            use_v4 = gr.Checkbox(label="Use v4", value=USE_V4)
             input = gr.Audio(sources=["upload", "microphone"], type="filepath", label="Audio")
 
         with gr.Column(scale=1):
@@ -40,8 +42,8 @@ with gr.Blocks(theme=theme) as demo:
             submit_btn = gr.Button("Submit", variant="primary")
 
 
-    submit_btn.click(fn=transcribe, inputs=[input], outputs=[output])
-    clear_btn.click(fn=clear,inputs=[], outputs=[input], queue=False,)
+    submit_btn.click(fn=transcribe, inputs=[input, use_v4], outputs=[output])
+    clear_btn.click(fn=clear,inputs=[], outputs=[input, use_v4], queue=False,)
 
 
 if __name__ == "__main__":

requirements.txt

@@ -1,5 +1,6 @@
 git+https://github.com/huggingface/transformers
 torch
+pyannote.audio
 yt-dlp
 gradio==4.20.0
 torchaudio==2.2.1

whisper.py

@@ -0,0 +1,205 @@
+from pyannote.audio import Pipeline
+from pydub import AudioSegment
+import os
+from transformers import WhisperForConditionalGeneration, WhisperProcessor
+import torchaudio
+import torch
+
+device = 0 if torch.cuda.is_available() else "cpu"
+torch_dtype = torch.float32
+
+
+MODEL_NAME = "openai/whisper-large-v3"
+model = WhisperForConditionalGeneration.from_pretrained(MODEL_NAME, torch_dtype=torch_dtype).to(device)
+processor = WhisperProcessor.from_pretrained(MODEL_NAME)
+pipeline_vad = Pipeline.from_pretrained("pyannote/voice-activity-detection", use_auth_token=os.environ.get("HF_TOKEN"))
+threshold = 15000  # adjust max duration threshold
+segments_dir = "."
+
+def clean_text(input_text):
+    
+    remove_chars = ['.', ',', ';', ':', '¿', '?', '«', '»', '-', '¡', '!', '@',
+                     '*', '{', '}', '[', ']', '=', '/', '\\', '&', '#', '…']
+
+    output_text = ''.join(char if char not in remove_chars else ' ' for char in input_text) #removing special chars
+    return (' '.join(output_text.split()).lower()) #remove extra spaces and return cleaned text
+
+def convert_forced_to_tokens(forced_decoder_ids):
+    forced_decoder_tokens = []
+    for i, (idx, token) in enumerate(forced_decoder_ids):
+        if token is not None:
+            forced_decoder_tokens.append([idx, processor.tokenizer.decode(token)])
+        else:
+            forced_decoder_tokens.append([idx, token])
+    return forced_decoder_tokens
+
+def generate_1st_chunk(audio):
+
+    input_audio, sample_rate = torchaudio.load(audio)
+    input_audio = torchaudio.transforms.Resample(sample_rate, 16000)(input_audio)
+    
+    input_speech = input_audio[0]
+
+    input_features = processor(input_speech, 
+                                    sampling_rate=16_000, 
+                                    return_tensors="pt", torch_dtype=torch_dtype).input_features.to(device)
+
+    forced_decoder_ids = []
+    forced_decoder_ids.append([1,50270]) #[1, '<|ca|>']
+    forced_decoder_ids.append([2,50262]) #[2, '<|es|>']
+    forced_decoder_ids.append([3,50360]) #[3, '<|transcribe|>']
+
+    forced_decoder_ids_modified = forced_decoder_ids
+
+    # we need to force these tokens
+    forced_decoder_ids = []
+
+    # now we need to append the prefix tokens (lang, task, timestamps)
+    offset = len(forced_decoder_ids)
+    for idx, token in forced_decoder_ids_modified:
+        forced_decoder_ids.append([idx + offset , token])
+    
+    model.generation_config.forced_decoder_ids = forced_decoder_ids
+
+    pred_ids = model.generate(input_features, 
+                                    return_timestamps=True,
+                                    max_new_tokens=128)
+    #exclude prompt from output
+    forced_decoder_tokens = convert_forced_to_tokens(forced_decoder_ids)
+    output = processor.decode(pred_ids[0][len(forced_decoder_tokens) + 1:], skip_special_tokens=True)
+    output_tokens = processor.batch_decode(pred_ids, skip_special_tokens=False)
+
+    return output[1:]
+
+def generate_from_2nd_chunk(audio, prev_prompt):
+
+    input_audio, sample_rate = torchaudio.load(audio)
+    input_audio = torchaudio.transforms.Resample(sample_rate, 16000)(input_audio)
+    
+    input_speech = input_audio[0]
+
+    input_features = processor(input_speech, 
+                                    sampling_rate=16_000, 
+                                    return_tensors="pt", torch_dtype=torch_dtype).input_features.to(device)
+    forced_decoder_ids = []
+
+    forced_decoder_ids.append([1,50270]) #[1, '<|ca|>']
+    forced_decoder_ids.append([2,50262]) #[2, '<|es|>']
+    forced_decoder_ids.append([3,50360]) #[3, '<|transcribe|>']
+
+    forced_decoder_ids_modified = forced_decoder_ids
+    idx = processor.tokenizer.all_special_tokens.index("<|startofprev|>")
+    forced_bos_token_id = processor.tokenizer.all_special_ids[idx]
+        
+    prompt_tokens = processor.tokenizer(prev_prompt, add_special_tokens=False).input_ids
+
+    # we need to force these tokens
+    forced_decoder_ids = []
+    for idx, token in enumerate(prompt_tokens):
+        # indexing starts from 1 for forced tokens (token at position 0 is the SOS token)
+        forced_decoder_ids.append([idx + 1, token])
+            
+    # now we add the SOS token at the end
+    offset = len(forced_decoder_ids)
+    forced_decoder_ids.append([offset + 1, model.generation_config.decoder_start_token_id])
+
+    # now we need to append the rest of the prefix tokens (lang, task, timestamps)
+    offset = len(forced_decoder_ids)
+    for idx, token in forced_decoder_ids_modified:
+        forced_decoder_ids.append([idx + offset , token])
+
+    model.generation_config.forced_decoder_ids = forced_decoder_ids
+
+    pred_ids = model.generate(input_features, 
+                                    return_timestamps=True,
+                                    max_new_tokens=128,
+                                    decoder_start_token_id=forced_bos_token_id)
+    #exclude prompt from output
+    forced_decoder_tokens = convert_forced_to_tokens(forced_decoder_ids)
+    output = processor.decode(pred_ids[0][len(forced_decoder_tokens) + 1:], skip_special_tokens=True)
+    output_tokens = processor.batch_decode(pred_ids, skip_special_tokens=False)
+    return output[1:]
+
+def processing_vad_v3(audio, output_vad, prev_prompt):
+    transcription_audio = ""
+    first_chunk = True
+    for speech in output_vad.get_timeline().support():
+        start, end = speech.start, speech.end
+        segment_audio = audio[start * 1000:end * 1000]
+        segment_audio.export(os.path.join(segments_dir, f"temp_segment.wav"), format="wav")
+        filename = os.path.join(segments_dir, f"temp_segment.wav")
+        if first_chunk:
+            output = generate_1st_chunk(filename)
+            first_chunk = False
+        else:
+            output = generate_from_2nd_chunk(filename, prev_prompt)
+
+        prev_prompt = output
+        transcription_audio = transcription_audio + " " + output
+
+    return transcription_audio
+
+
+def processing_vad_v4(audio, output_vad, threshold, max_duration, prev_prompt, concatenated_segment):
+    transcription_audio = ""
+    is_first_chunk = True
+    for speech in output_vad.get_timeline().support():
+        start, end = speech.start, speech.end
+        segment_duration = (end - start) * 1000
+        segment_audio = audio[start * 1000:end * 1000]
+
+        if max_duration + segment_duration < threshold:
+            concatenated_segment += audio[start * 1000:end * 1000]
+            max_duration += segment_duration
+        else:
+            if len(concatenated_segment) > 0:
+                temp_segment_path = os.path.join(segments_dir, f"temp_segment.wav")
+                concatenated_segment.export(temp_segment_path, format="wav")
+                
+                if is_first_chunk:
+                    output = generate_1st_chunk(temp_segment_path)
+                    is_first_chunk = False
+                else:
+                    output = generate_from_2nd_chunk(temp_segment_path, prev_prompt)
+
+                prev_prompt = output   
+                transcription_audio = transcription_audio + output
+
+                max_duration = segment_duration
+                concatenated_segment = segment_audio
+        
+    # Process any remaining audio in the concatenated_segment
+    if len(concatenated_segment) > 0:
+        temp_segment_path = os.path.join(segments_dir, f"temp_segment.wav")
+        concatenated_segment.export(temp_segment_path, format="wav")
+        
+        output = generate_from_2nd_chunk(temp_segment_path, prev_prompt)
+  
+        prev_prompt = output
+        transcription_audio = transcription_audio + output
+
+    return transcription_audio
+
+
+def generate(audio_path, use_v4):
+    #check audio lenght
+    audio = AudioSegment.from_wav(audio_path)
+    duration_seconds = len(audio) / 1000.0  
+
+    #apply VAD only if the duration is >30s
+    if duration_seconds >= 30:
+
+        output_vad = pipeline_vad(audio_path)
+        concatenated_segment = AudioSegment.empty()
+        max_duration = 0
+        prev_prompt = ""
+        if use_v4:
+            return processing_vad_v4(audio, output_vad, threshold, max_duration, prev_prompt, concatenated_segment)
+        else:
+            return  processing_vad_v3(audio, output_vad, prev_prompt)
+    else:
+        #if duraion is <30s, process directly with generate
+        return generate_1st_chunk(audio_path)
+        
+
+