import gradio as gr import torch from transformers import pipeline username = "ardneebwar" ## Complete your username model_id = f"{username}/distilhubert-finetuned-gtzan" device = "cuda:0" if torch.cuda.is_available() else "cpu" pipe = pipeline("audio-classification", model=model_id, device=device) # def predict_trunc(filepath): # preprocessed = pipe.preprocess(filepath) # truncated = pipe.feature_extractor.pad(preprocessed,truncation=True, max_length = 16_000*30) # model_outputs = pipe.forward(truncated) # outputs = pipe.postprocess(model_outputs) # return outputs 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) # preds = predict_trunc(filepath) outputs = {} for p in preds: outputs[p["label"]] = p["score"] return outputs title = "🎵 Music Genre Classifier" description = """ demo to showcase the music classification model that we just trained on the [GTZAN](https://huggingface.co/datasets/marsyas/gtzan) """ filenames = ['blues.00098.wav', "disco.00020.wav", "metal.00014.wav", "reggae.00021.wav", "rock.00058.wav"] filenames = [[f"./{f}"] for f in filenames] demo = gr.Interface( fn=classify_audio, inputs=gr.Audio(type="filepath"), outputs=gr.outputs.Label(), title=title, description=description, examples=filenames, ) demo.launch()