import gradio as gr
import numpy as np
from transformers import pipeline
import os

example_list = [["examples/" + example] for example in os.listdir("examples")]

classifier = pipeline('audio-classification', model='SamuelM0422/distilhubert-finetuned-gtzan')

title = 'Music Classification 🎙️'
description = 'A distilhubert model finetuned at gtzan dataset to classify music genres'

def predict(example):
  #print(type(example))
  example = {'array': np.array(example[1], dtype=np.float32), 'sampling_rate': example[0]}
  pred = classifier(example)
  return {p['label']: p['score'] for p in pred}

demo = gr.Interface(fn=predict, title=title, description=description,
                    inputs="audio", outputs="label", examples=example_list, flagging_mode='never')
demo.launch()