Upload main.py

main.py

@@ -0,0 +1,67 @@
+import gradio as gr
+from transformers import pipeline
+from transformers import MBartForConditionalGeneration, MBart50TokenizerFast
+from utils import lang_ids
+import nltk
+nltk.download('punkt')
+
+MODEL_NAME = "openai/whisper-medium"
+BATCH_SIZE = 8
+FILE_LIMIT_MB = 1000
+
+pipe = pipeline(
+    task="automatic-speech-recognition",
+    model=MODEL_NAME,
+    chunk_length_s=30,
+    device='cpu',
+)
+
+lang_list = list(lang_ids.keys())
+
+def translate_audio(inputs,target_language):
+    if inputs is None:
+        raise gr.Error("No audio file submitted! Please upload an audio file before submitting your request.")
+
+    text = pipe(inputs, batch_size=BATCH_SIZE, generate_kwargs={"task": "translate"}, return_timestamps=True)["text"]
+
+    target_lang = lang_ids[target_language]
+
+    if target_language == 'English':
+         return text
+
+    else:
+        model = MBartForConditionalGeneration.from_pretrained("sanjitaa/mbart-many-to-many")
+        tokenizer = MBart50TokenizerFast.from_pretrained("sanjitaa/mbart-many-to-many")
+
+        tokenizer.src_lang = "en_XX"
+        chunks = nltk.tokenize.sent_tokenize(text)
+        translated_text = ''
+
+        for segment in chunks:
+                encoded_chunk = tokenizer(segment, return_tensors="pt")
+                generated_tokens = model.generate(
+                             
+                     **encoded_chunk,
+                     forced_bos_token_id=tokenizer.lang_code_to_id[target_lang]
+                )
+                translated_chunk = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)
+                translated_text = translated_text + translated_chunk[0]
+        return translated_text
+
+inputs=[
+    gr.inputs.Audio(source="upload", type="filepath", label="Audio file"),
+    gr.Dropdown(lang_list, value="English", label="Target Language"),
+    ]
+description = "Audio translation"
+
+
+translation_interface = gr.Interface(
+    fn=translate_audio,
+    inputs= inputs,
+    outputs="text",
+    title="Speech Translation",
+    description= description
+)
+
+if __name__ == "__main__":
+    translation_interface.launch()