from transformers import AutoFeatureExtractor, AutoModelForAudioClassification
import gradio as gr

from pipeline import AudioPipeline


model_id = 'bjpietrzak/music_mind_distillhubert_gtzan_4e-5_WAdam_CosineCheguler'

feature_extractor = AutoFeatureExtractor.from_pretrained(
    model_id, do_normalize=True, return_attention_mask=True
)

model = AutoModelForAudioClassification.from_pretrained(model_id)

audio_pipeline = AudioPipeline(feature_extractor, model, top_k=7)


demo = gr.Interface(
    fn=audio_pipeline,
    inputs=[gr.Audio(type="filepath", label="Upload Audio")],
    outputs=gr.Label(num_top_classes=7),
    title="Music Mind",
)

demo.launch(debug=True)