import gradio as gr
from transformers import pipeline


model_id = "Teapack1/model_KWS"  # update with your model id
pipe = pipeline("audio-classification", model=model_id)


title = "Keyword Spotting Wav2Vec2"
description = "Gradio demo for finetuned Wav2Vec2 model on a custom dataset to perform keyword spotting task. Classes are scene 1, scene 2, scene 3, ambient, light on."

demo = gr.Blocks()

def classify_audio(audio):
    preds = pipe(audio)
    outputs = {}
    for p in preds:
        outputs[p["label"]] = p["score"]
    return outputs

    
mic_classify = gr.Interface(
    fn=classify_audio,
    inputs=gr.inputs.Audio(source="microphone", type="filepath", label="Record your audio"),
    outputs=gr.outputs.Label(),
    title=title,
    theme="huggingface",
    description=description,
    examples=[
        ["./scene3_329.wav",],
        ["./scene1_200.wav"],
        ["./light_422.wav"],
        ["./ambient_476.wav"],
        ],
    cache_examples=True,
)

file_classify = gr.Interface(
    fn=classify_audio,
    title=title,
    description=description,
    inputs=gr.inputs.Audio(source="upload", optional=True, label="Audio file", type="filepath"),
    theme="huggingface",
    outputs=gr.outputs.Label(),
)

# iface.test_examples(example_samples)


with demo:
    gr.TabbedInterface(
        [mic_classify, file_classify],
        ["Classify Microphone", "Classify Audio File"],
    )


demo.launch(debug=True, share=True)