Update app.py

app.py

@@ -4,6 +4,8 @@ import gradio as gr
 import yt_dlp as youtube_dl
 from transformers import pipeline
 from transformers.pipelines.audio_utils import ffmpeg_read
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
+from flores200_codes import flores_codes
 
 import tempfile
 import os
@@ -23,12 +25,32 @@ pipe = pipeline(
 )
 
 
+def load_translation_model():
+    model_name = 'facebook/nllb-200-distilled-1.3B'
+    model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    return model, tokenizer
+
+translation_model, translation_tokenizer = load_translation_model()
+
+
+def translate_text(text, source_language, target_language):
+    source_code = flores_codes[source_language]
+    target_code = flores_codes[target_language]
+
+    translator = pipeline('translation', model=translation_model, tokenizer=translation_tokenizer, src_lang=source_code, tgt_lang=target_code)
+    output = translator(text, max_length=400)
+    return output[0]['translation_text']
+
+
+
 def transcribe(inputs, task):
     if inputs is None:
         raise gr.Error("No audio file submitted! Please upload or record an audio file before submitting your request.")
 
     text = pipe(inputs, batch_size=BATCH_SIZE, generate_kwargs={"task": task}, return_timestamps=True)["text"]
-    return  text
+    translated_text = translate_text(text, source_language, target_language)
+    return  text, translated_text
 
 
 def _return_yt_html_embed(yt_url):
@@ -88,6 +110,9 @@ def yt_transcribe(yt_url, task, max_filesize=75.0):
     return html_embed_str, text
 
 
+lang_codes = list(flores_codes.keys())
+
+
 demo = gr.Blocks()
 
 mf_transcribe = gr.Interface(
@@ -95,8 +120,10 @@ mf_transcribe = gr.Interface(
     inputs=[
         gr.inputs.Audio(source="microphone", type="filepath", optional=True),
         gr.inputs.Radio(["transcribe", "translate"], label="Task", default="transcribe"),
+        gr.inputs.Dropdown(lang_codes, default='English', label='Source Language'),
+        gr.inputs.Dropdown(lang_codes, default='French', label='Target Language'),
     ],
-    outputs="text",
+    outputs=["text", "text"],
     layout="horizontal",
     theme="huggingface",
     title="Whisper Large V2: Transcribe Audio",
@@ -113,6 +140,8 @@ file_transcribe = gr.Interface(
     inputs=[
         gr.inputs.Audio(source="upload", type="filepath", optional=True, label="Audio file"),
         gr.inputs.Radio(["transcribe", "translate"], label="Task", default="transcribe"),
+        gr.inputs.Dropdown(lang_codes, default='English', label='Source'),
+        gr.inputs.Dropdown(lang_codes, default='French', label='Target'),
     ],
     outputs="text",
     layout="horizontal",
@@ -130,7 +159,9 @@ yt_transcribe = gr.Interface(
     fn=yt_transcribe,
     inputs=[
         gr.inputs.Textbox(lines=1, placeholder="Paste the URL to a YouTube video here", label="YouTube URL"),
-        gr.inputs.Radio(["transcribe", "translate"], label="Task", default="transcribe")
+        gr.inputs.Radio(["transcribe", "translate"], label="Task", default="transcribe"),
+        gr.inputs.Dropdown(lang_codes, default='English', label='Source Language'),
+        gr.inputs.Dropdown(lang_codes, default='French', label='Target Language'),
     ],
     outputs=["html", "text"],
     layout="horizontal",