Upload app.py

app.py

@@ -11,7 +11,7 @@ import gradio as gr
 import numpy as np
 import torch
 
-from transformers import pipeline, VitsModel, VitsTokenizer
+from transformers import pipeline, VitsModel, VitsTokenizer, FSMTForConditionalGeneration, FSMTTokenizer
 
 device = "cuda:0" if torch.cuda.is_available() else "cpu"
 
@@ -19,7 +19,11 @@ device = "cuda:0" if torch.cuda.is_available() else "cpu"
 asr_pipe = pipeline("automatic-speech-recognition", model="asapp/sew-d-tiny-100k-ft-ls100h", device=device)
 
 #eng text to rus text translation
-translation_pipe = pipeline("translation", model="facebook/wmt19-en-ru")
+mname = "facebook/wmt19-en-ru"
+tokenizer = FSMTTokenizer.from_pretrained(mname)
+model = FSMTForConditionalGeneration.from_pretrained(mname)
+
+#translation_pipe = pipeline("translation", model="facebook/wmt19-en-ru")
 
 #rus text to rus speech transformation
 vits_model = VitsModel.from_pretrained("facebook/mms-tts-rus")
@@ -30,8 +34,12 @@ def transform_audio_to_speech_en(audio):
     return outputs["text"]
 
 def translator(text):
-  translated_text = translation_pipe(text)
-  return translated_text[0]['translation_text']
+  input_ids = tokenizer.encode(text, return_tensors="pt")
+  outputs = model.generate(input_ids)
+  decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)
+  return decoded
+  #translated_text = translation_pipe(text)
+  #return translated_text[0]['translation_text']
 
 def synthesise(translated_text):
     translated_text = translator(translated_text)