Update app.py

app.py

@@ -1,5 +1,6 @@
 import torch
 import numpy as np
+import soundfile as sf
 from transformers import pipeline
 from transformers import BarkModel
 from transformers import AutoProcessor
@@ -9,26 +10,29 @@ device="cpu"
 pipe = pipeline(
     "automatic-speech-recognition", model="openai/whisper-large-v2", device=device
 )
+label = pipeline("audio-classification", model="facebook/mms-lid-126", device=device)
 processor = AutoProcessor.from_pretrained("suno/bark")
 model = BarkModel.from_pretrained("suno/bark")
 model = model.to(device)
 synthesised_rate = model.generation_config.sample_rate
 
-def translate(audio):
+def translate(audio_file):
+    audio, sampling_rate = sf.read(audio_file)
     outputs = pipe(audio, max_new_tokens=256, generate_kwargs={"task": "transcribe","language":"chinese"})
-    return outputs["text"]
+    language_prediction = label({"array": audio, "sampling_rate": sampling_rate})
+    label_outputs = {}
+    for pred in language_prediction:
+        label_outputs[pred["label"]] = pred["score"]
+    return outputs["text"],label_outputs
 def synthesise(text_prompt,voice_preset="v2/zh_speaker_1"):
     inputs = processor(text_prompt, voice_preset=voice_preset)
     speech_output = model.generate(**inputs.to(device),pad_token_id=10000)  
     return speech_output
 def speech_to_speech_translation(audio,voice_preset="v2/zh_speaker_1"):
-    translated_text = translate(audio)
+    translated_text, label_outputs= translate(audio)
     synthesised_speech = synthesise(translated_text,voice_preset)
     synthesised_speech = (synthesised_speech.numpy() * 32767).astype(np.int16)
-    return synthesised_rate , synthesised_speech ,translated_text
-def speech_to_speech_translation_fix(audio,voice_preset="v2/zh_speaker_1"):
-    synthesised_rate,synthesised_speech,translated_text = speech_to_speech_translation(audio,voice_preset)
-    return (synthesised_rate,synthesised_speech.T),translated_text
+    return (synthesised_rate , synthesised_speech.T),translated_text,label_outputs
 
 title = "Multilanguage to Chinese(mandarin) Cascaded STST"
 description = """
@@ -51,22 +55,24 @@ import gradio as gr
 
 demo = gr.Blocks()
 file_transcribe = gr.Interface(
-    fn=speech_to_speech_translation_fix,
+    fn=speech_to_speech_translation,
     inputs=gr.Audio(source="upload", type="filepath"),
     outputs=[
         gr.Audio(label="Generated Speech", type="numpy"),
         gr.Text(label="Transcription"),
+        gr.Label(label="Language prediction"),
     ],
     title=title,
     description=description,
     examples=examples,
 )
 mic_transcribe = gr.Interface(
-    fn=speech_to_speech_translation_fix,
+    fn=speech_to_speech_translation,
     inputs=gr.Audio(source="microphone", type="filepath"),
     outputs=[
         gr.Audio(label="Generated Speech", type="numpy"),
         gr.Text(label="Transcription"),
+        gr.Label(label="Language prediction"),
     ],
     title=title,
     description=description,