from transformers import pipeline
import librosa
import gradio as gr


config = {
    "sampling_rate": 16000,
    "model": "bjpietrzak/distilhubert-gtzan-20-5e-5",
    "top_k": 7}

def audio_pipeline(file_path: str) -> dict[str, float]:
    y, _ = librosa.load(file_path, sr=config['sampling_rate'])
    out = pipe(y, top_k=config['top_k'])
    return {clas['label']: clas['score'] for clas in out}


pipe = pipeline("audio-classification", model=config['model'])

demo = gr.Interface(
    fn=audio_pipeline,
    inputs=[gr.Audio(type="filepath")],
    outputs=gr.Label(num_top_classes=7),
    title="Music Mind Audio Classification",
    description="Upload an .mp3 or .ogg audio file "
    "to classify the content using a pre-trained model.")

demo.launch(debug=True)