Update app.py

app.py

@@ -28,7 +28,7 @@ def whisper_demo(input_audio, model_id):
     pipe.model.config.forced_decoder_ids = pipe.tokenizer.get_decoder_prompt_ids(language='en', task="transcribe")
     output_text = pipe(input_audio)['text']
     return output_text
-    
+
     
 def translate_to_english(prompt, lang_model_id, base_lang):
     if base_lang == "English":
@@ -73,21 +73,24 @@ def biogpt_audio(
     input_audio: str,
     biogpt_model_id: str,
     whisper_model_id: str,
-    max_length: str,
-    num_return_sequences: int
+    base_lang: str,
+    lang_model_id: str,
 ):
     en_prompt = whisper_demo(input_audio=input_audio, model_id=whisper_model_id)
     generator = pipeline("text-generation", model=biogpt_model_id, device="cuda:0")
-    output = generator(en_prompt, max_length=max_length, num_return_sequences=num_return_sequences, do_sample=True)
-    output_dict = {}
-    for i in range(num_return_sequences):
-        output_dict[str(i+1)] = output[i]['generated_text']
-
-    output_text = ""
-    for i in range(num_return_sequences):
-        output_text += f'{output_dict[str(i+1)]}\n\n'
-    
-    return en_prompt, output_text, output_text
+    output = generator(en_prompt, max_length=250, num_return_sequences=1, do_sample=True)
+    if base_lang == "English":
+        output_text = output
+        
+    else: 
+        output_text = text_to_text_generation(
+            prompt=output,
+            model_id=lang_model_id,
+            device='cuda:0',
+            target_lang=lang_ids[base_lang]
+        )
+            
+    return en_prompt, output, output_text
 
 examples = [
     ["COVID-19 is", biogpt_model_list[0], lang_model_list[1], "English"]
@@ -95,7 +98,7 @@ examples = [
 
 app = gr.Blocks()
 with app:
-    gr.Markdown("# **<h4 align='center'>Whisper + M2M100 + BioGPT: Generative Pre-trained Transformer for Biomedical Text Generation and Mining</h4>**")
+    gr.Markdown("# **<h4 align='center'>Whisper + M2M100 + BioGPT: Generative Pre-trained Transformer for Biomedical Text Generation and Mining<h4>**")
     gr.Markdown(
         """
         <p style='text-align: center'>
@@ -109,12 +112,16 @@ with app:
             with gr.Tab("Text"):
                 input_text = gr.Textbox(lines=3, value="COVID-19 is", label="Text")
                 input_text_button = gr.Button(value="Predict")
-                input_biogpt_model =gr.Dropdown(choices=biogpt_model_list, value=biogpt_model_list[0], label='BioGpt Model')
-                input_m2m100_model =gr.Dropdown(choices=lang_model_list,  value=lang_model_list[1], label='Language Model')
+                input_biogpt_model = gr.Dropdown(choices=biogpt_model_list, value=biogpt_model_list[0], label='BioGpt Model')
+                input_m2m100_model = gr.Dropdown(choices=lang_model_list,  value=lang_model_list[1], label='Language Model')
                 input_base_lang = gr.Dropdown(lang_list, value="English", label="Base Language")
             
             with gr.Tab("Audio"):
-                input_audio = gr.Microphone(label='Audio')
+                input_audio = gr.Audio(source="microphone", type="filepath")
+                input_biogpt_model = gr.Dropdown(choices=biogpt_model_list, value=biogpt_model_list[0], label='BioGpt Model')
+                input_whisper_model = gr.Dropdown(choices=whisper_model_list,  value=whisper_model_list[0], label='Audio Model')
+                input_base_lang = gr.Dropdown(lang_list, value="English", label="Base Language")
+                input_m2m100_model = gr.Dropdown(choices=lang_model_list,  value=lang_model_list[1], label='Language Model')
                 input_audio_button = gr.Button(value="Predict")    
 
         with gr.Column():
@@ -124,6 +131,6 @@ with app:
                 
     gr.Examples(examples, inputs=[input_text, input_biogpt_model, input_m2m100_model,input_base_lang], outputs=[prompt_text, output_text, translated_text], fn=biogpt_text, cache_examples=False)
     input_text_button.click(biogpt_text, inputs=[input_text, input_biogpt_model, input_m2m100_model,input_base_lang], outputs=[prompt_text, output_text, translated_text])
-    input_audio_button.click(biogpt_audio, inputs=[input_audio, input_biogpt_model,input_m2m100_model,input_base_lang], outputs=[prompt_text, output_text, translated_text])
+    input_audio_button.click(biogpt_audio, inputs=[input_audio, input_biogpt_model,input_whisper_model,input_base_lang], outputs=[prompt_text, output_text, translated_text])
     
 app.launch()