Update fo Russian speech app.py

app.py

@@ -13,31 +13,25 @@ model_wav2vec = 'openai/whisper-small'  #'voidful/wav2vec2-xlsr-multilingual-56'
 
 asr_pipe = pipeline("automatic-speech-recognition", model=model_wav2vec, device=device)
 
-# from speech to text
+# load speech-to-text checkpoint
 def translate_audio(audio):
-    outputs = asr_pipe(audio, max_new_tokens=256, 
-                       generate_kwargs={"task": "translate"})
+    outputs = asr_pipe(audio, max_new_tokens=256, generate_kwargs={"task": "translate"})
     return outputs["text"]
 
-# translation
-def translate_text(text, from_language, target_language):  #to English -mul en, to Russian - en ru
-    model_name = f'Helsinki-NLP/opus-mt-{from_language}-{target_language}'
-    tokenizer = MarianTokenizer.from_pretrained(model_name)
-    model = MarianMTModel.from_pretrained(model_name)
-
-    inputs = tokenizer.encode(text, return_tensors="pt")
-    outputs = model.generate(inputs, num_beams=4, max_length=50, early_stopping=True)
-    translated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return translated_text
+# translation to Russian
+def translate_text(text):  
+    # to English - mul en, to Russian - en ru
+    model_mul_en = pipeline("translation", model = "Helsinki-NLP/opus-mt-mul-en")
+    model_en_ru = pipeline("translation", model = "Helsinki-NLP/opus-mt-en-ru")
+    translated_text = model_en_ru(model_mul_en(text)[0]['translation_text'])
+    return translated_text[0]['translation_text']
 
 # load text-to-speech checkpoint
-#model = pipeline("text-to-speech", model="voxxer/speecht5_finetuned_commonvoice_ru_translit")
 model = VitsModel.from_pretrained("voxxer/speecht5_finetuned_commonvoice_ru_translit")
 tokenizer = VitsTokenizer.from_pretrained("facebook/mms-tts-rus")
 
 def synthesise(text):
-    translated_text = translate_text(text, 'mul', 'en')
-    translated_text = translate_text(translate_text, 'en', 'ru')
+    translated_text = translate_text(text)
     inputs = tokenizer(translated_text, return_tensors="pt")
     input_ids = inputs["input_ids"]
     with torch.no_grad():
@@ -55,7 +49,7 @@ def speech_to_speech_translation(audio):
 
 title = "Cascaded STST"
 description = """
-* В начале происходит распознавание речи с помощью модели voidful/wav2vec2-xlsr-multilingual-56 и на выходе получается текст на любом из 56 языков.
+* В начале происходит распознавание речи с помощью модели openai/whisper-small.
 * Затем полученный текст переводится сначала на английский с помощью Helsinki-NLP/opus-mt-mul-en, а потом на русский с помощью Helsinki-NLP/opus-mt-en-ru
 * На последнем шаге полученный текст озвучивается с помощью fine-tune-говой версии microsoft/speecht5_tts - voxxer/speecht5_finetuned_commonvoice_ru_translit
 Demo for cascaded speech-to-speech translation (STST), mapping from source speech in any language to target speech in Russian. Demo uses facebook/mms-tts-rus model for text-to-speech: