import gradio as gr
from transformers import pipeline

model_name = "juliensimon/wav2vec2-conformer-rel-pos-large-finetuned-speech-commands"

p = pipeline("audio-classification", model=model_name)

def process(file):
        pred = p(file)
        labels = dict()
        for l in pred:
            labels[l['label']]=l['score']
        return labels

# Gradio inputs
mic = gr.inputs.Audio(source='microphone', type='filepath', label='Speech input', optional=True)

# Gradio outputs
keyword = gr.outputs.Label(num_top_classes=3)

description = "This Space showcases a wav2vec2-conformer-rel-pos-large model fine-tuned for audio classification on the speech_commands dataset. \n \n It can spot one of the following keywords: 'Yes', 'No', 'Up', 'Down', 'Left', 'Right', 'On', 'Off', 'Stop', 'Go', 'Zero', 'One', 'Two', 'Three', 'Four', 'Five', 'Six', 'Seven', 'Eight', 'Nine', 'Bed', 'Bird', 'Cat', 'Dog', 'Happy', 'House', 'Marvin', 'Sheila', 'Tree', 'Wow', 'Backward', 'Forward', 'Follow', 'Learn', 'Visual'."

iface = gr.Interface(
	theme='huggingface',
	description=description,
	fn=process,
        layout='horizontal',
	inputs=[mic],
	outputs=[keyword],
	examples=[
		['backward16k.wav'],
		['happy16k.wav'],
		['marvin16k.wav'],
		['seven16k.wav'],
		['stop16k.wav'],
		['up16k.wav'],
	],
    allow_flagging=False
)
iface.launch()