Update whisper.py

whisper.py

@@ -9,33 +9,12 @@ device = 0 if torch.cuda.is_available() else "cpu"
 torch_dtype = torch.float32
 
 HF_TOKEN = os.getenv("HF_TOKEN")
-#MODEL_NAME = "openai/whisper-large-v3"
 MODEL_NAME = "projecte-aina/whisper-large-v3-ca-es-synth-cs"
 model = WhisperForConditionalGeneration.from_pretrained(MODEL_NAME, torch_dtype=torch_dtype,token=HF_TOKEN).to(device)
 processor = WhisperProcessor.from_pretrained(MODEL_NAME)
-pipeline_vad = Pipeline.from_pretrained("./pyannote/config.yaml") 
-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):
 
+def generate(audio_path):
     input_audio, sample_rate = torchaudio.load(audio)
     input_audio = torchaudio.transforms.Resample(sample_rate, 16000)(input_audio)
     
@@ -44,161 +23,11 @@ def generate_1st_chunk(audio):
     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)
-
-    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)
-    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]
-        filename = os.path.join(segments_dir, f"temp_segment.wav")
-        segment_audio.export(filename, format="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)
-        
 
-   
+    output = processor.batch_decode(pred_ids, skip_special_tokens=True)
+    line = output[0]
+    return line