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

pipe = pipeline("audio-classification", model="TheDuyx/distilhubert-bass-classifier5")

def classify_audio(filepath):
    audio, sampling_rate = librosa.load(filepath, sr=16_000)
    preds = pipe(audio)
    outputs = {}
    for p in preds:
        outputs[p["label"]] = p["score"]
    return outputs

demo = gr.Interface(
    fn=classify_audio,
    inputs=gr.Audio(type="filepath"),
    outputs="label",
    examples=[["brass.wav"], ["growl.wav"], ["808.wav"], ["acid.wav"], ["slap.wav"]],
)

demo.launch()