Upload 2 files

app.py

@@ -1,58 +1,39 @@
-import spaces
-import torch
 import gradio as gr
-from AinaTheme import theme
-from transformers import pipeline
-
-MODEL_NAME = "projecte-aina/whisper-large-v3-ca-es-synth-cs"
-BATCH_SIZE = 8
-device = 0 if torch.cuda.is_available() else "cpu"
-
-pipe = pipeline(
-    task="automatic-speech-recognition",
-    model=MODEL_NAME,
-    chunk_length_s=30,
-    device=device,
-)
-
-@spaces.GPU
-def transcribe(inputs):
+from whisper import generate
+from AinaTheme import theme  
+
+USE_V5 = False
+
+def transcribe(inputs, model_version):
     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")
-    text = pipe(inputs, batch_size=BATCH_SIZE, generate_kwargs={"task": "transcribe"}, return_timestamps=True)["text"]
-    return text
 
+    use_v5 = model_version == "v0.5"
+    return generate(audio_path=inputs, use_v5=use_v5)
 
 description_string = "Transcripció automàtica de micròfon o de fitxers d'àudio.\n Aquest demostrador s'ha desenvolupat per"\
-              " comprovar els models de reconeixement de parla per a móbils. Per ara utilitza el checkpoint "\
-              f"[{MODEL_NAME}](https://huggingface.co/{MODEL_NAME}) i la llibreria de 🤗 Transformers per a la transcripció."
-
+              " comprovar els models de reconeixement de parla per a móbils."
 
 def clear():
-     return (
-          None
-     )
-
+    return None, "v1.0"
 
-with gr.Blocks(theme=theme) as demo:
+with gr.Blocks() as demo:
     gr.Markdown(description_string)
     with gr.Row():
         with gr.Column(scale=1):
-            #input = gr.Audio(sources=["upload", "microphone"], type="filepath", label="Audio")
-            input = gr.Audio(sources=["upload"], type="filepath", label="Audio")
+            model_version = gr.Dropdown(label="Model Version", choices=["v1.0", "v0.5"], value="v1.0")
+            input = gr.Audio(sources=["upload", "microphone"], type="filepath", label="Audio")
 
         with gr.Column(scale=1):
             output = gr.Textbox(label="Output", lines=8)
-    
-    with gr.Row(variant="panel"):
-            clear_btn = gr.Button("Clear")
-            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,)
+    with gr.Row(variant="panel"):
+        clear_btn = gr.Button("Clear")
+        submit_btn = gr.Button("Submit", variant="primary")
 
+    submit_btn.click(fn=transcribe, inputs=[input, model_version], outputs=[output])
+    clear_btn.click(fn=clear, inputs=[], outputs=[input, model_version], queue=False)
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()

whisper.py

@@ -1,34 +1,240 @@
-import os
 from pyannote.audio import Pipeline
 from pydub import AudioSegment
+import os
 from transformers import WhisperForConditionalGeneration, WhisperProcessor
 import torchaudio
 import torch
+import re
+from transformers import pipeline
+
 
 device = 0 if torch.cuda.is_available() else "cpu"
 torch_dtype = torch.float32
 
-HF_TOKEN = os.getenv("HF_TOKEN")
-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)
+
+MODEL_NAME = "openai/whisper-large-v3"
+CKPT = "projecte-aina/whisper-large-v3-tiny-caesar"
+BATCH_SIZE = 1
+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 = 10000 
+segments_dir = "."
+
+pipe = pipeline(
+    task="automatic-speech-recognition",
+    model=CKPT,
+    chunk_length_s=30,
+    device=device
+    ) 
+
+def post_process_transcription(example_transcription, max_repeats=1):
+    segments = re.findall(r'.+?[.,?]', example_transcription)
+
+    seen = set()
+    unique_segments = []
+    for segment in segments:
+        if segment not in seen:
+            unique_segments.append(segment)
+            seen.add(segment)
+
+    final_string = ''.join(unique_segments)
+
+    tokens = re.findall(r'\b\w+\b[.,!?]?', final_string)
+
+    cleaned_tokens = []
+    repetition_count = 0
+    previous_token = None
+
+    for token in tokens:
+        reduced_token = re.sub(r"(\w{1,3})(\1{2,})", "", token)
+
+        if reduced_token == previous_token:
+            repetition_count += 1
+            if repetition_count <= max_repeats:
+                cleaned_tokens.append(reduced_token)
+        else:
+            repetition_count = 1
+            cleaned_tokens.append(reduced_token)
+
+        previous_token = reduced_token
 
+    cleaned_transcription = " ".join(cleaned_tokens)
+    cleaned_transcription = re.sub(r'\s+', ' ', cleaned_transcription).strip()
 
-def generate(audio_path):
-    input_audio, sample_rate = torchaudio.load(audio_path)
+    return cleaned_transcription
+
+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_attention_mask=True,
+                                    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 = "Antes de 'digui'm', '112'. 112, digui'm. Hola, puc parlar en castellà? Sí, digui, diga. Sí, mire: a veces al abrir la puerta de mi piso tengo una persona ahí. Vale, avisamos a la Guàrdia Urbana, ¿de acuerdo? Vale, perfecto. Gracias. Gracias. Buen día."
+    prompt_tokens = processor.tokenizer(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 generate_2nd_chuk(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
+    idx = processor.tokenizer.all_special_tokens.index("<|startofprev|>")
+    forced_bos_token_id = processor.tokenizer.all_special_ids[idx]
+        
+    prompt = "112, digui'm. Hola, puc parlar en castellà? Sí, digui, diga. Sí, mire: a veces al abrir la puerta de mi piso tengo una persona ahí. Vale, avisamos a la Guàrdia Urbana, ¿de acuerdo? Vale, perfecto. Gracias. Gracias. Buen día."  
+    prompt_tokens = processor.tokenizer(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)
+                                    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_threshold(audio, output_vad, threshold, max_duration, 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_2nd_chuk(temp_segment_path)
+                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_2nd_chuk(temp_segment_path)
+        transcription_audio = transcription_audio + output
+
+    return(transcription_audio)
+
+def format_audio(audio):
+    input_audio, sample_rate = torchaudio.load(audio)
+    resampler = torchaudio.transforms.Resample(sample_rate, 16000)
+    input_audio = resampler(input_audio)
+    input_audio = input_audio.squeeze().numpy()
+    return(input_audio)
+
+def transcribe_pipeline(audio, task):
+    text = pipe(audio, batch_size=BATCH_SIZE, generate_kwargs={"task": task}, return_timestamps=True)["text"]
+    return text
+
+def generate(audio_path, use_v5):
+    audio = AudioSegment.from_wav(audio_path)
+
+    output_vad = pipeline_vad(audio_path)
+    concatenated_segment = AudioSegment.empty()
+    max_duration = 0
+
+    if use_v5:
+        output = processing_vad_threshold(audio, output_vad, threshold, max_duration, concatenated_segment)
+    else:  
+        task = "transcribe"
+        output = transcribe_pipeline(format_audio(audio), task)
+        
+    clean_output = post_process_transcription(output, max_repeats=1)
+
+    return clean_output 
 
-    output = processor.batch_decode(pred_ids, skip_special_tokens=True)
-    line = output[0]
-    return line
+