different translation 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, translate
-# from translation import get_translation
+from speech2text import init_model_trans, transcribe
+from translation import translate
 from resources import audit_elapsedtime, set_start
 import subprocess
 
@@ -29,9 +29,10 @@ def main ():
         start_loading = set_start()
         st.audio(wav_audio_data, format='audio/wav')
         original = transcribe(wav_audio_data, s2t)
+        st.write(f"Original: {original}")
         print("translating audio...")
-        translation = translate(original, s2t)
-        st.write(f"Original: {original}/nTranscription: {translation}")
+        translation = translate(original)
+        st.write(f"Transcription: {translation}")
 
     #     if text is not None and ner is not None:    
     #         st.write('Entities: ', get_entity_labels(model=ner, text=text))

translation.py

@@ -1,5 +1,6 @@
 from transformers import MarianMTModel, MarianTokenizer
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
+from resources import set_start, audit_elapsedtime
 
 def get_model_name(languageCode: str) -> str:
     match languageCode:
@@ -8,12 +9,15 @@ def get_model_name(languageCode: str) -> str:
     
     return model_name
 
-def get_translation(text_to_translate: str, languageCode: str) -> str:
-        
+def translate(text_to_translate: str) -> str:
+    
+    start = set_start()
+    print("Initiating translation model...")
     tokenizer = AutoTokenizer.from_pretrained("unicamp-dl/translation-pt-en-t5")
     model = AutoModelForSeq2SeqLM.from_pretrained("unicamp-dl/translation-pt-en-t5")
     pten_pipeline = pipeline('text2text-generation', model=model, tokenizer=tokenizer)
     translated_text = pten_pipeline(text_to_translate)
 
+    audit_elapsedtime(function="Finished translation", start=start)
     print("Translated text:", translated_text)
     return translated_text