added translation task in whisper model

app.py

@@ -2,8 +2,8 @@ import streamlit as st
 from st_audiorec import st_audiorec
 
 from nameder import init_model_ner, get_entity_labels
-from speech2text import init_model_trans, transcribe
-from translation import get_translation
+from speech2text import init_model_trans, transcribe, translate
+# from translation import get_translation
 from resources import audit_elapsedtime, set_start
 import subprocess
 
@@ -12,25 +12,26 @@ def main ():
     print(f"Running main")
 
     #print(subprocess.Popen('pip freeze > requirements_hug.txt', shell=True))
-    text = "Tenho uma proposta para a Caixa Geral de Depositos, para 3 consultores outsystems, 300 euros por dia e um periodo de seis meses."
-    st.write(text)
-    traducao = get_translation(text_to_translate=text, languageCode="pt")
-    st.write(traducao)
-    # s2t = init_model_trans()
+    # text = "Tenho uma proposta para a Caixa Geral de Depositos, para 3 consultores outsystems, 300 euros por dia e um periodo de seis meses."
+    # st.write(text)
+    # traducao = get_translation(text_to_translate=text, languageCode="pt")
+    # st.write(traducao)
+    s2t = init_model_trans()
     # ner = init_model_ner() #async
 
-    # print("Rendering UI...")
-    # start_render = set_start()
-    # wav_audio_data = st_audiorec()
-    # audit_elapsedtime(function="Rendering UI", start=start_render)
+    print("Rendering UI...")
+    start_render = set_start()
+    wav_audio_data = st_audiorec()
+    audit_elapsedtime(function="Rendering UI", start=start_render)
 
-    # if wav_audio_data is not None and s2t is not None:
-    #     print("Loading data...")
-    #     start_loading = set_start()
-    #     st.audio(wav_audio_data, format='audio/wav')
-    #     text = transcribe(wav_audio_data, s2t)
-    #     print("translating audio...")
-    #     translation = get_translation("pt")
+    if wav_audio_data is not None and s2t is not None:
+        print("Loading data...")
+        start_loading = set_start()
+        st.audio(wav_audio_data, format='audio/wav')
+        original = transcribe(wav_audio_data, s2t)
+        print("translating audio...")
+        translation = translate(original)
+        st.write(f"Original: {original}/nTranscription: {translation}")
 
     #     if text is not None and ner is not None:    
     #         st.write('Entities: ', get_entity_labels(model=ner, text=text))

speech2text.py

@@ -14,7 +14,7 @@ def init_model_trans ():
     model_id = "openai/whisper-large-v3"
 
     model = AutoModelForSpeechSeq2Seq.from_pretrained(
-        model_id, torch_dtype=torch_dtype, low_cpu_mem_usage=False, use_safetensors=True
+        model_id, torch_dtype=torch_dtype, low_cpu_mem_usage=True, use_safetensors=True
     )
     model.to(device)
 
@@ -41,9 +41,27 @@ def transcribe (audio_sample: bytes, pipe) -> str:
     start = set_start()
     # dataset = load_dataset("distil-whisper/librispeech_long", "clean", split="validation")
     # sample = dataset[0]["audio"]
+    
+    #result = pipe(audio_sample)
     result = pipe(audio_sample)
+
     audit_elapsedtime(function="Transcription", start=start)
     print("transcription result",result)
 
+    #st.write('trancription: ', result["text"])
+    return result["text"]
+
+def translate (audio_sample: bytes, pipe) -> str:
+    print("Initiating Translation...")
+    start = set_start()
+    # dataset = load_dataset("distil-whisper/librispeech_long", "clean", split="validation")
+    # sample = dataset[0]["audio"]
+    
+    #result = pipe(audio_sample)
+    result = pipe(audio_sample, generate_kwargs={"task": "translate"})
+
+    audit_elapsedtime(function="Translation", start=start)
+    print("Translation result",result)
+
     #st.write('trancription: ', result["text"])
     return result["text"]