Update app.py

app.py

@@ -1,6 +1,13 @@
 import gradio as gr
 from transformers import pipeline
 
+pretrained_model: str = "facebook/m2m100_1.2B"
+cache_dir: str = "models/"
+
+tokenizer = M2M100Tokenizer.from_pretrained(pretrained_model, cache_dir=cache_dir)
+translation_model = M2M100ForConditionalGeneration.from_pretrained(
+        pretrained_model, cache_dir=cache_dir)
+
 transcription = pipeline("automatic-speech-recognition", model= "openai/whisper-base")
 clasification = pipeline(
     "audio-classification",
@@ -117,20 +124,33 @@ lang_id = {
     "Chinese": "zh",
     "Zulu": "zu",
 }
-
+def translation(source_lang, target_lang):
+    src_lang = lang_id[source_lang]
+    trg_lang = lang_id[target_lang]
+    tokenizer.src_lang = src_lang
+    with torch.no_grad():
+        encoded_input = tokenizer(user_input, return_tensors="pt").to(device)
+        generated_tokens = model.generate(
+        **encoded_input, forced_bos_token_id=tokenizer.get_lang_id(trg_lang))
+        translated_text = tokenizer.batch_decode(
+        generated_tokens, skip_special_tokens=True)[0]
+    return translated_text
 demo = gr.Blocks()
 
 with demo:
   gr.Markdown("Speech analyzer")
   audio = gr.Audio(type="filepath", label = "Upload a file")
+  text0 = gr.Textbox()
   text = gr.Textbox()
   source_lang = gr.Dropdown(label="Source lang", choices=list(lang_id.keys()))
+  target_lang = gr.Dropdown(label="target lang", choices=list(lang_id.keys()))
+   
   #gr.Examples(examples = list(lang_id.keys()),
    #             inputs=[
     #                source_lang])
   b1 = gr.Button("convert to text")
-  #b3 = gr.Button("select audio")
-  #b3.Dropdown(inputs = source_lang)
+  b3 = gr.Button("translate")
+  b3.Dropdown(translation, input = text0, output = text)
   b1.click(audio_a_text, inputs=audio, outputs=text)
 
   b2 = gr.Button("Classification of speech")