Update app.py

app.py

@@ -1,5 +1,6 @@
 import gradio as gr 
 from transformers import pipeline, Wav2Vec2ProcessorWithLM, Wav2Vec2ForCTC
+from transformers import T5ForConditionalGeneration, T5Tokenizer
 import os
 import soundfile as sf
 from pyannote.audio import Pipeline
@@ -9,7 +10,7 @@ from pydub.playback import play
 from datetime import datetime, timedelta
 import time
 from silero_vad import load_silero_vad, read_audio, get_speech_timestamps
-from googletrans import Translator
+#from googletrans import Translator
 
 
 HF_TOKEN = os.environ.get("HF_TOKEN")
@@ -26,8 +27,17 @@ model = load_silero_vad()
 pipeline_dia = Pipeline.from_pretrained("pyannote/speaker-diarization-3.1",
                                     use_auth_token=HF_TOKEN)
 
+device = 'cpu' #or 'cuda' for translate on cpu
 
-translator = Translator()
+model_name_tr = 'utrobinmv/t5_translate_en_ru_zh_small_1024'
+model_tr = T5ForConditionalGeneration.from_pretrained(model_name_tr)
+model_tr.to(device)
+tokenizer_tr = T5Tokenizer.from_pretrained(model_name_tr)
+
+prefix = 'translate to en: '
+
+
+# translator = Translator()
 temp_path = "temp.wav"
 
 def preprocess(audio_path):
@@ -66,7 +76,13 @@ def fast_transcribe(diarise, how_diarise, translate, audio):
                       lines.append(f"{i+1}\n{start_time_srt} --> {end_time_srt}\n[{label}] {trans}\n")
                     else:
                       print("TRANSLATION STARTED")
-                      trans_eng = translator.translate('trans', src='ru', dest="en").text
+                      src_text = prefix + trans
+                      # translate Russian to Eng
+                      input_ids = tokenizer_tr(src_text, return_tensors="pt")
+                      generated_tokens = model_tr.generate(**input_ids.to(device))
+
+                      trans_eng = tokenizer_tr.batch_decode(generated_tokens, skip_special_tokens=True)
+                      #trans_eng = translator.translate(trans, src='ru', dest="en").text
                       print(f"TRANSLATION ENDED RESULT {trans_eng}")
                       lines.append(f"{i+1}\n{start_time_srt} --> {end_time_srt}\n[{label}] {trans}\n[{label}] {trans_eng}\n")
                     print("RECOGNISING ENDED")
@@ -93,7 +109,13 @@ def fast_transcribe(diarise, how_diarise, translate, audio):
                   lines.append(f"{i+1}\n{start_time_srt} --> {end_time_srt}\n[{trans}\n")
                 else:
                   print("TRANSLATION STARTED")
-                  trans_eng = translator.translate(trans, src='ru', dest="en").text
+                  src_text = prefix + trans
+                  # translate Russian to Eng
+                  input_ids = tokenizer_tr(src_text, return_tensors="pt")
+                  generated_tokens = model_tr.generate(**input_ids.to(device))
+
+                  trans_eng = tokenizer_tr.batch_decode(generated_tokens, skip_special_tokens=True)
+                  #trans_eng = translator.translate(trans, src='ru', dest="en").text
                   print(f"TRANSLATION ENDED RESULT {trans_eng}")
                   lines.append(f"{i+1}\n{start_time_srt} --> {end_time_srt}\n{trans}\n{trans_eng}\n")