Finishing

app.py

@@ -1,49 +1,24 @@
-import torch
-import torchaudio
 from transformers import AutoFeatureExtractor, AutoModelForAudioClassification
 import gradio as gr
-import json
 
+from pipeline import AudioPipeline
 
-config = json.load(open("config.json"))
+
+model_id = 'bjpietrzak/music_mind_distillhubert_gtzan_4e-5_WAdam_CosineCheguler'
 
 feature_extractor = AutoFeatureExtractor.from_pretrained(
-    config['extractor_path'], do_normalize=True, return_attention_mask=True
+    model_id, do_normalize=True, return_attention_mask=True
 )
 
-model = AutoModelForAudioClassification.from_pretrained(
-    config['model_path']
-)
+model = AutoModelForAudioClassification.from_pretrained(model_id)
+
+audio_pipeline = AudioPipeline(feature_extractor, model, top_k=7)
 
-def audio_pipeline(audio_file):
-    if isinstance(audio_file, str):
-        waveform, sample_rate = torchaudio.load(audio_file)
-    else:
-        waveform, sample_rate = torchaudio.load(audio_file.name)
-    waveform = waveform.mean(dim=0)
-    if sample_rate != feature_extractor.sampling_rate:
-        transform = torchaudio.transforms.Resample(
-            orig_freq=sample_rate,
-            new_freq=feature_extractor.sampling_rate)
-        waveform = transform(waveform)
-    inputs = feature_extractor(waveform,
-        sampling_rate=feature_extractor.sampling_rate,
-        return_tensors="pt",
-        padding=True)
-    with torch.no_grad():
-        logits = model(**inputs).logits
-    probs = torch.nn.functional.softmax(logits, dim=-1)[0]
-    
-    top_probs, top_ids = torch.topk(probs, config['top_k'])
-    top_labels = [model.config.id2label[idx.item()] for idx in top_ids]
-    
-    results = {label: prob.item() for label, prob in zip(top_labels, top_probs)}
-    return results
 
 demo = gr.Interface(
     fn=audio_pipeline,
     inputs=[gr.Audio(type="filepath", label="Upload Audio")],
-    outputs=gr.Label(num_top_classes=config['top_k']),
+    outputs=gr.Label(num_top_classes=7),
     title="Music Mind",
 )
 

pipeline.py

@@ -0,0 +1,33 @@
+import torch
+import torchaudio
+
+
+class AudioPipeline:
+    def __init__(self, feature_extractor, model, top_k=5):
+        self.fe = feature_extractor
+        self.model = model
+        self.top_k = top_k
+
+    def __call__(self, audio_file):
+        if isinstance(audio_file, str):
+            waveform, sample_rate = torchaudio.load(audio_file)
+        else:
+            waveform, sample_rate = torchaudio.load(audio_file.name)
+        waveform = waveform.mean(dim=0)
+        if sample_rate != self.fe.sampling_rate:
+            transform = torchaudio.transforms.Resample(
+                orig_freq=sample_rate,
+                new_freq=self.fe.sampling_rate)
+            waveform = transform(waveform)
+        inputs = self.fe(waveform,
+            sampling_rate=self.fe.sampling_rate,
+            return_tensors="pt",
+            padding=True)
+        with torch.no_grad():
+            logits = self.model(**inputs).logits
+        probs = torch.nn.functional.softmax(logits, dim=-1)[0]
+        
+        top_probs, top_ids = torch.topk(probs, self.top_k)
+        top_labels = [self.model.config.id2label[idx.item()] for idx in top_ids]
+        
+        return {label: prob.item() for label, prob in zip(top_labels, top_probs)}