Use 3 stage cascade for better results

app.py

@@ -4,14 +4,18 @@ import numpy as np
 import torch
 
 from transformers import VitsModel, VitsTokenizer, pipeline
+from transformers import M2M100ForConditionalGeneration
+from tokenization_small100 import SMALL100Tokenizer
 
 
 device = "cuda:0" if torch.cuda.is_available() else "cpu"
 
-target_language = "french"
+target_language = "fr"
 
 # load speech translation checkpoint
 asr_pipe = pipeline("automatic-speech-recognition", model="bofenghuang/whisper-small-cv11-french", device=device)
+translation_model = M2M100ForConditionalGeneration.from_pretrained("alirezamsh/small100", device=device)
+translation_tokenizer = SMALL100Tokenizer.from_pretrained("alirezamsh/small100", tgt_lang=target_language)
 
 # load text-to-speech checkpoint
 model = VitsModel.from_pretrained("facebook/mms-tts-fra")
@@ -19,8 +23,11 @@ tokenizer = VitsTokenizer.from_pretrained("facebook/mms-tts-fra")
 
 
 def translate(audio):
-    outputs = asr_pipe(audio, max_new_tokens=256, generate_kwargs={"task": "transcribe", "language": target_language})
-    return outputs["text"]
+    outputs = asr_pipe(audio, max_new_tokens=256, generate_kwargs={"task": "translate"})
+    eng_text = outputs["text"]
+    encoded_eng_text = translation_tokenizer(eng_text, return_tensors="pt")
+    generated_tokens = translation_model(**encoded_eng_text)
+    return translation_tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)
 
 
 def synthesise(text):