Update app.py

app.py

@@ -3,11 +3,68 @@ import torch
 from transformers import AutoProcessor, VoxtralForConditionalGeneration
 import spaces
 
+#### Functions
+
+@spaces.GPU
+def process_transcript(language: str, audio_path: str) -> str:
+    """Process the audio file to return its transcription.
+
+    Args:
+        language: The language of the audio.
+        audio_path: The path to the audio file.
+
+    Returns:
+        The transcribed text of the audio.
+    """
+
+    if audio_path is None:
+        return "Please provide some input audio: either upload an audio file or use the microphone."
+    else:
+        id_language = dict_languages[language]
+        inputs = processor.apply_transcrition_request(language=id_language, audio=audio_path, model_id=model_name)
+        inputs = inputs.to(device, dtype=torch.bfloat16)
+        outputs = model.generate(**inputs, max_new_tokens=MAX_TOKENS)
+        decoded_outputs = processor.batch_decode(outputs[:, inputs.input_ids.shape[1]:], skip_special_tokens=True)
+
+        return decoded_outputs[0]
+###
+
+def process_translate(language: str, audio_path: str) -> str:
+    conversation = [
+        {
+            "role": "user",
+            "content": [
+                {
+                    "type": "audio",
+                    "path": "+audio_path,
+                },
+                {"type": "text", "text": "Translate this in "+language},
+            ],
+        }
+    ]
+    
+    inputs = processor.apply_chat_template(conversation)
+    inputs = inputs.to(device, dtype=torch.bfloat16)
+    
+    outputs = model.generate(**inputs, max_new_tokens=MAX_TOKENS)
+    decoded_outputs = processor.batch_decode(outputs[:, inputs.input_ids.shape[1]:], skip_special_tokens=True)
+
+    return decoded_outputs
+
+
+def disable_buttons():
+    return gr.update(interactive=False), gr.update(interactive=False)
+
+def enable_buttons():
+    return gr.update(interactive=True), gr.update(interactive=True)
+###
+
+### Initializations
+
 MAX_TOKENS = 32000
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 print(f"*** Device: {device}")
-
 model_name = 'mistralai/Voxtral-Mini-3B-2507'
 
 processor = AutoProcessor.from_pretrained(model_name)
@@ -24,28 +81,13 @@ dict_languages = {"English": "en",
                   "Dutch": "nl",
                   "Hindi": "hi"}
 
-@spaces.GPU
-def process_transcript(language, audio_path):
-    """Process audio with selected Voxtral model and return the generated response"""
-
-    if audio_path is None:
-        return "Please provide some input audio: either upload an audio file or use the microphone."
-    else:
-        id_language = dict_languages[language]
-        inputs = processor.apply_transcrition_request(language=id_language, audio=audio_path, model_id=model_name)
-        inputs = inputs.to(device, dtype=torch.bfloat16)
-        outputs = model.generate(**inputs, max_new_tokens=MAX_TOKENS)
-        decoded_outputs = processor.batch_decode(outputs[:, inputs.input_ids.shape[1]:], skip_special_tokens=True)
-
-        return decoded_outputs[0]
-
 
+#### Gradio interface
 
-# Define Gradio interface
 with gr.Blocks(title="Transcription") as transcript:
-    gr.Markdown("# Audio Transcription")
-    gr.Markdown("#### Choose the language of the audio and the model, then set an audio file to get its transcription.")
-    gr.Markdown("#### **(Voxtral handles audios up to 30 minutes for transcription)**")
+    gr.Markdown("# Voxtral Mini Evaluation")
+    gr.Markdown("#### Choose the language of the audio and set an audio file to process it.")
+    gr.Markdown("##### *(Voxtral handles audios up to 30 minutes for transcription)*")
 
     with gr.Row():
         with gr.Column():
@@ -56,7 +98,8 @@ with gr.Blocks(title="Transcription") as transcript:
             )
 
         with gr.Column():
-            sel_audio = gr.Audio(sources=["microphone", "upload"], type="filepath", label="Upload an audio file or record via microphone:")
+            sel_audio = gr.Audio(sources=["upload", "microphone"], type="filepath", 
+                                 label="Upload an audio file, record via microphone, or select a demo file:")
 
             example = [["mapo_tofu.mp3"]]
             gr.Examples(
@@ -70,16 +113,55 @@ with gr.Blocks(title="Transcription") as transcript:
     
     with gr.Row():
         with gr.Column():
-            submit_transcript = gr.Button("Extract Transcription", variant="primary")
-            text_transcript = gr.Textbox(label="Generated Response", lines=10)
+            submit_transcript = gr.Button("Extract transcription", variant="primary")
+            text_transcript = gr.Textbox(label="Generated transcription", lines=10)
 
+        with gr.Column():
+            sel_translate_language = gr.Dropdown(
+                choices=list(dict_languages.keys()),
+                value="English",
+                label="Select the language for translation:"
+            )
+
+            submit_translate = gr.Button("Translate audio file", variant="primary")
+            text_translate = gr.Textbox(label="Generated translation", lines=10)
+
+        with gr.Column():
+            submit_chat = gr.Button("Ask audio file", variant="primary")
+            text_chat = gr.Textbox(label="Model answer", lines=10)
             
+### Processing
+    
+    # Transcription
     submit_transcript.click(
+        disable_buttons,
+        outputs=[submit_transcript, submit_translate, submit_chat],
+        trigger_mode="once",
+    ).then(
         fn=process_transcript,
         inputs=[sel_language, sel_audio],
         outputs=text_transcript
+    ).then(
+        enable_buttons,
+        outputs=[submit_transcript, submit_translate, submit_chat],
     )
 
-# Launch the app
+    # Translation
+    submit_transcript.click(
+        disable_buttons,
+        outputs=[submit_transcript, submit_translate, submit_chat],
+        trigger_mode="once",
+    ).then(
+        fn=process_transcript,
+        inputs=[sel_language, sel_audio],
+        outputs=text_transcript
+    ).then(
+        enable_buttons,
+        outputs=[submit_transcript, submit_translate, submit_chat],
+    )
+
+
+### Launch the app
+
 if __name__ == "__main__":
-    transcript.launch(share=True)
+    audio.launch()