from transformers import pipeline
import gradio
from gradio import Interface, Audio, Label, Number

username = 'bvallegc' ## Complete your username
model_id = f"{username}/distilhubert-finetuned-gtzan"
pipe = pipeline("audio-classification", model=model_id)

def classify_audio(filepath):
    """
    Goes from
    [{'score': 0.8339303731918335, 'label': 'country'},
  {'score': 0.11914275586605072, 'label': 'rock'},]

   to
   {"country":  0.8339303731918335, "rock":0.11914275586605072}
  """
    preds = pipe(filepath)
    classification = [{"label": p["label"], "score": p["score"]} for p in preds]
    label = classification[0]["label"]
    number = classification[0]["score"]
    return label, number

example_audio_files = [
    "Freedom.mp3",
    "In The Forest.mp3",
    "Summer Mood.mp3",
]

examples = [{"filepath": path} for path in example_audio_files]

interface_options = {
    "title": "Music Genre Classification",
    "description": "The audio classifier for those who are the best and only want and require the best",
    # Audio from validation file
    "allow_flagging": "never"
}

demo = Interface(
    fn=classify_audio, inputs= Audio(type="filepath", examples=examples), outputs=[Label(), Number()], **interface_options
)
demo.launch(debug=False)