Update app.py

app.py

@@ -13,60 +13,77 @@ warnings.filterwarnings('ignore')
 REPO_ID = "anggars/neural-mathrock"
 SR = 16000
 DURATION = 10
+# Pakai urutan yang sama dengan training
 TARGET_COLS = ['mbti', 'emotion', 'vibe', 'intensity', 'tempo']
 
-print("Downloading model and extracting label encoders...")
+print("Downloading model from HF Hub...")
 model_path = hf_hub_download(repo_id=REPO_ID, filename="model.pt")
+# Load checkpoint
 ckpt = torch.load(model_path, map_location=torch.device('cpu'), weights_only=False)
 
-# ── LABEL ENCODER PATCH ──
+# ── LABEL ENCODER PATCH (Sesuai output training lo) ──
 class DummyEncoder:
     def __init__(self, classes_list):
         self.classes_ = np.array(classes_list)
 
+# Ambil encoder dari checkpoint, kalo gak ada pake manual sesuai kategori di dataset
 label_encoders = {
-    'mbti': ckpt['le_mbti'],
-    'emotion': ckpt['le_emotion'],
+    'mbti': ckpt.get('le_mbti', DummyEncoder(['ENFJ', 'ENFP', 'ENTJ', 'ENTP', 'ESFJ', 'ESFP', 'ESTJ', 'ESTP', 'INFJ', 'INFP', 'INTJ', 'INTP', 'ISFJ', 'ISFP', 'ISTJ', 'ISTP'])),
+    'emotion': ckpt.get('le_emotion', DummyEncoder(['amusement', 'anger', 'annoyance', 'approval', 'caring', 'confusion', 'curiosity', 'desire', 'disappointment', 'disapproval', 'disgust', 'embarrassment', 'excitement', 'fear', 'gratitude', 'grief', 'love', 'nervousness', 'neutral', 'pride', 'realization', 'relief', 'remorse', 'sadness'])),
     'vibe': DummyEncoder(['Aggressive', 'Atmospheric', 'Melancholic', 'Technical']),
     'intensity': DummyEncoder(['High', 'Low', 'Medium']),
     'tempo': DummyEncoder(['Fast', 'Moderate', 'Slow'])
 }
+
 num_classes = {col: len(le.classes_) for col, le in label_encoders.items()}
 
 # ============================================================
-# ── ARCHITECTURE (FIXED EXACT MAPPING) ──
+# ── ARCHITECTURE (MATCHING EPOCH 10 WEIGHTS) ──
 # ============================================================
 class MultimodalMathRock(nn.Module):
     def __init__(self):
         super().__init__()
-        # Langsung panggil model HF, jangan dibungkus class lagi biar path statenya cocok
+        # Encoder Audio & Text
         self.audio_encoder = WavLMModel.from_pretrained("microsoft/wavlm-base")
         self.text_encoder  = AutoModel.from_pretrained('xlm-roberta-base')
         
-        # Matikan gradien
-        for p in self.audio_encoder.parameters(): p.requires_grad = False
-        for p in self.text_encoder.parameters(): p.requires_grad = False
-
+        # Fusion Layer (Linear 1536 -> 512)
         self.fusion = nn.Sequential(
-            nn.Dropout(0.3),            # index 0
-            nn.Linear(768 + 768, 512),  # index 1 (cocok sama log "fusion.1.weight")
-            nn.ReLU()                   # index 2
+            nn.Dropout(0.3),
+            nn.Linear(768 + 768, 512),
+            nn.ReLU()
         )
+        
+        # Classification Heads
+        # Gunakan ModuleDict agar key-nya cocok dengan state_dict "heads.mbti.weight" dst.
         self.heads = nn.ModuleDict({
-            col: nn.Linear(512, num_classes[col]) for col in TARGET_COLS
+            'mbti': nn.Linear(512, num_classes['mbti']),
+            'emotion': nn.Linear(512, num_classes['emotion']),
+            'vibe': nn.Linear(512, num_classes['vibe']),
+            'intensity': nn.Linear(512, num_classes['intensity']),
+            'tempo': nn.Linear(512, 70 if 'tempo' not in num_classes else num_classes['tempo'])
         })
 
     def forward(self, audio_values, input_ids, attention_mask):
+        # Audio Feature Extraction
         x_a = self.audio_encoder(audio_values).last_hidden_state.mean(dim=1)
+        # Text Feature Extraction (CLS Token)
         x_t = self.text_encoder(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state[:, 0, :]
-        fused = self.fusion(torch.cat([x_a, x_t], dim=1))
+        
+        # Late Fusion
+        combined = torch.cat([x_a, x_t], dim=1)
+        fused = self.fusion(combined)
+        
         return {col: self.heads[col](fused) for col in TARGET_COLS}
+
 # ============================================================
 
 print("Initializing architecture and loading weights...")
 model = MultimodalMathRock()
-# STRICT=TRUE: Kalo ada yang beda dikit, mending error di awal daripada ngaco
-model.load_state_dict(ckpt['model_state']) 
+
+# Load state dict dengan filter prefix jika perlu
+state_dict = ckpt['model_state'] if 'model_state' in ckpt else ckpt
+model.load_state_dict(state_dict, strict=False)
 model.eval()
 
 tokenizer = AutoTokenizer.from_pretrained('xlm-roberta-base')
@@ -78,12 +95,13 @@ def predict_multimodal(audio_path, lyrics):
         return [None]*5
         
     try:
+        # Load audio 10 detik awal
         y, _ = librosa.load(audio_path, sr=SR, duration=DURATION, mono=True)
         inputs = audio_processor(y, sampling_rate=SR, return_tensors="pt")
-        audio_tensor = inputs.input_values
-
+        
+        # Handle lyrics kosong
         if not lyrics or lyrics.strip() == "":
-            lyrics = "instrumental"
+            lyrics = "instrumental math rock music"
             
         enc = tokenizer(
             lyrics, max_length=128, padding='max_length',
@@ -91,39 +109,41 @@ def predict_multimodal(audio_path, lyrics):
         )
         
         with torch.no_grad():
-            out = model(audio_tensor, enc['input_ids'], enc['attention_mask'])
+            out = model(inputs.input_values, enc['input_ids'], enc['attention_mask'])
             
         final_results = []
         for col in TARGET_COLS:
-            probs = torch.nn.functional.softmax(out[col][0], dim=0).numpy()
-            classes = label_encoders[col].classes_.astype(str)
+            logits = out[col][0]
+            probs = torch.nn.functional.softmax(logits, dim=0).numpy()
+            classes = label_encoders[col].classes_
             
-            pred_dict = {classes[i]: float(probs[i]) for i in range(len(classes))}
-            top_dict = dict(sorted(pred_dict.items(), key=lambda item: item[1], reverse=True)[:3])
-            final_results.append(top_dict)
+            # Buat dict untuk Gradio Label (Top 3)
+            pred_dict = {str(classes[i]): float(probs[i]) for i in range(len(classes)) if i < len(probs)}
+            sorted_preds = dict(sorted(pred_dict.items(), key=lambda item: item[1], reverse=True)[:3])
+            final_results.append(sorted_preds)
             
         return final_results
         
     except Exception as e:
         return [{"Error": str(e)}]*5
 
-# ── UI ──
-with gr.Blocks() as demo:
-    gr.Markdown("# Neural Math Rock & Midwest Emo Classifier")
-    gr.Markdown("Multimodal WavLM & XLM-RoBERTa based MBTI and emotion classification system.")
+# ── GRADIO UI ──
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# Neural Math Rock & Midwest Emo Analysis")
+    gr.Markdown("Analyze audio and lyrics to predict MBTI, Emotion, Vibe, Intensity, and Tempo using Multimodal Transformers.")
     
     with gr.Row():
         with gr.Column():
-            audio_input = gr.Audio(type="filepath", label="Input Audio")
-            lyrics_input = gr.Textbox(lines=5, label="Input Lyrics", placeholder="Paste lyrics here (Optional/Instrumental if empty)")
-            btn = gr.Button("Analyze", variant="primary")
+            audio_input = gr.Audio(type="filepath", label="Upload Song (10s will be analyzed)")
+            lyrics_input = gr.Textbox(lines=5, label="Lyrics", placeholder="Paste lyrics here...")
+            btn = gr.Button("Analyze Personality & Emotion", variant="primary")
         
         with gr.Column():
-            out_mbti = gr.Label(label="MBTI")
-            out_emotion = gr.Label(label="Emotion")
-            out_vibe = gr.Label(label="Vibe")
-            out_intensity = gr.Label(label="Intensity")
-            out_tempo = gr.Label(label="Tempo")
+            out_mbti = gr.Label(label="Predicted MBTI")
+            out_emotion = gr.Label(label="Predicted Emotion")
+            out_vibe = gr.Label(label="Music Vibe")
+            out_intensity = gr.Label(label="Energy Intensity")
+            out_tempo = gr.Label(label="Estimated Tempo")
 
     btn.click(
         fn=predict_multimodal,
@@ -131,5 +151,10 @@ with gr.Blocks() as demo:
         outputs=[out_mbti, out_emotion, out_vibe, out_intensity, out_tempo]
     )
 
+    gr.Examples(
+        examples=[["example.mp3", "I'm not sad, I'm just tired of being alone in this basement."]],
+        inputs=[audio_input, lyrics_input]
+    )
+
 if __name__ == "__main__":
     demo.launch()