import gradio as gr
from funasr import AutoModel

# Initialize the FunASR model with the specified components
model = AutoModel(
    model="paraformer-zh", model_revision="v2.0.4",
    vad_model="fsmn-vad", vad_model_revision="v2.0.4",
    punc_model="ct-punc-c", punc_model_revision="v2.0.4",
    # Uncomment the next line to enable speaker verification/diarization
    # spk_model="cam++", spk_model_revision="v2.0.2",
)

def transcribe(audio_file):
    # Processing the input audio file
    result = model.generate(
        input=audio_file.name,
        batch_size_s=300, 
        hotword='魔搭'  # This is an example keyword; replace or remove as necessary
    )
    return result

# Gradio interface setup
interface = gr.Interface(
    fn=transcribe,
    inputs=gr.inputs.Audio(source="microphone", type="filepath", label="Upload your audio in Mandarin"),
    outputs="text",
    title="FunASR Speech Recognition",
    description="This Gradio app uses the paraformer-zh model for speech recognition with additional features like VAD and punctuation restoration."
)

# Run the Gradio app
if __name__ == "__main__":
    interface.launch()