Update app.py

app.py

@@ -1,10 +1,20 @@
 import gradio as gr
-from faster_whisper import WhisperModel
+from transformers import pipeline
 from transformers import MBartForConditionalGeneration, MBart50TokenizerFast
 from utils import lang_ids
+import nltk
+nltk.download('punkt')
 
-model_size = "medium"
-ts_model = WhisperModel(model_size, device = "cpu", compute_type = "int8")
+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())
 
@@ -12,25 +22,23 @@ 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.")
 
-    segments, _ = ts_model.transcribe(inputs, task="translate")
+    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':
-        lst = ''
-        for segment in segments:
-             lst = lst + segment.text
-        return lst
+         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 segments:
-                encoded_chunk = tokenizer(segment.text, return_tensors="pt")
+        for segment in chunks:
+                encoded_chunk = tokenizer(segment, return_tensors="pt")
                 generated_tokens = model.generate(
                              
                      **encoded_chunk,
@@ -40,20 +48,19 @@ def translate_audio(inputs,target_language):
                 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=[
-        gr.inputs.Audio(source="upload", type="filepath", label="Audio file"),
-        gr.Dropdown(lang_list, value="English", label="Target Language"),
-    ],
+    inputs= inputs,
     outputs="text",
-    layout="horizontal",
-    theme="huggingface",
-    title="Translate Audio to English",
-    description=(
-        "Translate audio inputs to English using the"
-    ),
-    allow_flagging="never",
+    title="Speech Translation",
+    description= description
 )
 
 if __name__ == "__main__":