Update app.py

app.py

@@ -1,53 +1,22 @@
 import gradio as gr
-import torch
 from transformers import pipeline
+import torch
 
 model_id = "ntu-spml/distilhubert"
-device = "cuda:0" if torch.cuda.is_available() else "cpu"
+device = 0 if torch.cuda.is_available() else -1
 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):
     import time
-    start_time = time.time()
-    
-    # Assuming `pipe` is your model pipeline for inference
+    start = time.time()
     preds = pipe(filepath)
-    
-    outputs = {}
-    for p in preds:
-        outputs[p["label"]] = p["score"]
-    
-    end_time = time.time()
-    prediction_time = end_time - start_time
-    
-    return outputs, prediction_time
-
+    result = {p["label"]: round(p["score"], 3) for p in preds}
+    return result, round(time.time() - start, 2)
 
-title = "🎵 Music Genre Classifier"
-description = """
-Music Genre Classifier model (Fine-tuned "ntu-spml/distilhubert") Dataset: [GTZAN](https://huggingface.co/datasets/marsyas/gtzan)
-"""
-
-filenames = ['rock-it-21275.mp3']
-filenames = [f"./{f}" for f in filenames]
-
-demo = gr.Interface(
+gr.Interface(
     fn=classify_audio,
-    inputs=gr.Audio(type="filepath"),
-    outputs=[gr.Label(), gr.Number(label="Prediction time (s)")],
-    title=title,
-    description=description,
-)
-
-
-
-demo.launch()
+    inputs=gr.Audio(type="filepath", label="Upload Audio"),
+    outputs=[gr.Label(label="Top Genres"), gr.Number(label="Time (s)")],
+    title="🎵 Music Genre Classifier",
+    description="Classifies the genre of uploaded audio using DistilHuBERT fine-tuned on GTZAN."
+).launch()