Update app.py

app.py

@@ -1,9 +1,10 @@
 import torch
 import torch.nn as nn
+import torch.nn.functional as F
 import librosa
 import numpy as np
 import gradio as gr
-from transformers import XLMRobertaModel, XLMRobertaTokenizer, WavLMModel
+from transformers import XLMRobertaForSequenceClassification, XLMRobertaTokenizer, WavLMModel
 from huggingface_hub import hf_hub_download
 import yt_dlp
 import os
@@ -16,62 +17,46 @@ import warnings
 warnings.filterwarnings('ignore')
 
 # -- CONFIGURATION --
-REPO_ID = "anggars/neural-mathrock"
+REPO_AUDIO = "anggars/neural-mathrock"
+REPO_TEXT_MBTI = "anggars/xlm-mbti"
+REPO_TEXT_EMO = "anggars/xlm-emotion"
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-TARGET_COLS = ['mbti', 'emotion', 'vibe', 'intensity', 'tempo']
 GENIUS_TOKEN = os.environ.get("GENIUS_TOKEN", "z2XGBWXalGUtAdC1qxxXBxUnK1ZuoHPkCu5eP9q-fed-DW1uCJ3NSFpHemk3Unmg")
 
-print("Fetching model weights...")
-model_path = hf_hub_download(repo_id=REPO_ID, filename="model.pt")
-ckpt = torch.load(model_path, map_location=DEVICE, weights_only=False)
+# -- GLOBAL LABELS --
+MBTI_LABELS = sorted(["INTJ", "INTP", "ENTJ", "ENTP", "INFJ", "INFP", "ENFJ", "ENFP", "ISTJ", "ISFJ", "ESTJ", "ESFJ", "ISTP", "ISFP", "ESTP", "ESFP"])
+EMO_LABELS = sorted(['admiration', 'amusement', 'anger', 'annoyance', 'approval', 'caring', 'confusion', 'curiosity', 'desire', 'disappointment', 'disapproval', 'disgust', 'embarrassment', 'excitement', 'fear', 'gratitude', 'grief', 'joy', 'love', 'nervousness', 'optimism', 'pride', 'realization', 'relief', 'remorse', 'sadness', 'surprise', 'neutral'])
+VIBE_LABELS = sorted(["Aggressive", "Atmospheric", "Melancholic", "Technical"])
+INTENSITY_LABELS = sorted(["High", "Low", "Medium"])
 
-le_mbti, le_emotion, le_vibe, le_intensity, le_tempo = ckpt['le_mbti'], ckpt['le_emotion'], ckpt['le_vibe'], ckpt['le_intensity'], ckpt['le_tempo']
-
-# -- ARCHITECTURE --
-class HybridMultimodalModel(nn.Module):
+# -- AUDIO ARCHITECTURE (neural-mathrock) --
+class NeuralMathRockAudio(nn.Module):
     def __init__(self):
         super().__init__()
-        self.text_model = XLMRobertaModel.from_pretrained('anggars/xlm-mbti')
-        self.audio_model = WavLMModel.from_pretrained('microsoft/wavlm-base')
-        self.audio_proj = nn.Linear(768, 256)
-        self.text_gate = nn.Sequential(nn.Linear(768, 768), nn.Sigmoid())
-        self.audio_gate = nn.Sequential(nn.Linear(256, 256), nn.Sigmoid())
-        self.fusion = nn.Sequential(
-            nn.Linear(1024, 512),
-            nn.BatchNorm1d(512),
-            nn.ReLU(),
-            nn.Dropout(0.4),
-        )
-        self.head_mbti = nn.Linear(512, len(le_mbti.classes_))
-        self.head_emotion = nn.Linear(512, len(le_emotion.classes_))
-        self.head_vibe = nn.Linear(512, len(le_vibe.classes_))
-        self.head_intensity = nn.Linear(512, len(le_intensity.classes_))
-        self.head_tempo = nn.Linear(512, len(le_tempo.classes_))
-
-    def forward(self, input_ids, attention_mask, audio_values, text_missing=False):
-        text_out = self.text_model(input_ids=input_ids, attention_mask=attention_mask)
-        text_feat = text_out.pooler_output
-        
-        # Jika instrumental, matikan representasi teks
-        if text_missing:
-            text_feat = torch.zeros_like(text_feat)
-            
-        audio_out = self.audio_model(audio_values).last_hidden_state
-        audio_feat = self.audio_proj(audio_out.mean(dim=1))
-        
-        # Gating mechanism
-        gated_text = text_feat * self.text_gate(text_feat)
-        gated_audio = audio_feat * self.audio_gate(audio_feat)
-        
-        # Fusion
-        fused = self.fusion(torch.cat([gated_text, gated_audio], dim=-1))
-        
-        return {col: getattr(self, f'head_{col}')(fused) for col in TARGET_COLS}
+        self.wavlm = WavLMModel.from_pretrained("microsoft/wavlm-base-plus")
+        self.proj = nn.Sequential(nn.Linear(768, 256), nn.BatchNorm1d(256), nn.GELU(), nn.Dropout(0.4))
+        self.mbti_head = nn.Linear(256, len(MBTI_LABELS))
+        self.emo_head = nn.Linear(256, len(EMO_LABELS))
+        self.vibe_head = nn.Linear(256, len(VIBE_LABELS))
+        self.intensity_head = nn.Linear(256, len(INTENSITY_LABELS))
+        self.tempo_head = nn.Linear(256, 1)
 
-model = HybridMultimodalModel().to(DEVICE)
-model.load_state_dict(ckpt['model_state'], strict=False)
-model.eval()
-tokenizer = XLMRobertaTokenizer.from_pretrained('anggars/xlm-mbti')
+    def forward(self, iv, aam):
+        h = self.wavlm(iv, attention_mask=aam).last_hidden_state.mean(dim=1)
+        f = self.proj(h)
+        return self.mbti_head(f), self.emo_head(f), self.vibe_head(f), self.intensity_head(f), self.tempo_head(f).squeeze(-1)
+
+# -- MODEL INITIALIZATION --
+print("Initializing Hybrid Ensemble Models...")
+ckpt_path = hf_hub_download(repo_id=REPO_AUDIO, filename="model.pt")
+ckpt = torch.load(ckpt_path, map_location=DEVICE, weights_only=False)
+audio_model = NeuralMathRockAudio().to(DEVICE)
+audio_model.load_state_dict(ckpt['model'], strict=False)
+audio_model.eval()
+
+tokenizer = XLMRobertaTokenizer.from_pretrained(REPO_TEXT_MBTI)
+text_mbti_model = XLMRobertaForSequenceClassification.from_pretrained(REPO_TEXT_MBTI).to(DEVICE).eval()
+text_emo_model = XLMRobertaForSequenceClassification.from_pretrained(REPO_TEXT_EMO).to(DEVICE).eval()
 
 # -- UTILITIES --
 def search_and_fetch(query):
@@ -80,7 +65,6 @@ def search_and_fetch(query):
         search = VideosSearch(query, limit=1)
         res = search.result()
         if not res['result']: return None, "No results."
-        
         video_url = res['result'][0]['link']
         temp_fn = "temp_audio_file"
         ydl_opts = {
@@ -107,69 +91,81 @@ def search_and_fetch(query):
 def analyze_track(audio_path, lyrics_input):
     if not audio_path: return [{"Error": "No audio"}] * 5
     try:
-        is_inst = not lyrics_input or str(lyrics_input).strip() == ""
-        text = str(lyrics_input).strip() if not is_inst else "[INSTRUMENTAL]"
-        enc = tokenizer(text, truncation=True, padding='max_length', max_length=128, return_tensors='pt').to(DEVICE)
-        
         wav, sr = librosa.load(audio_path, sr=16000)
-        tempo_bpm, _ = librosa.beat.beat_track(y=wav, sr=sr)
-
         chunk_len = 16000 * 15
         chunks = [wav[i:i + chunk_len] for i in range(0, len(wav), chunk_len) if len(wav[i:i+chunk_len]) >= 16000]
-        all_logits = {col: [] for col in TARGET_COLS}
+        
+        a_mbti, a_emo, a_vibe, a_int, a_tmp = [], [], [], [], []
         
         with torch.no_grad():
             for chunk in chunks:
                 if len(chunk) < chunk_len: chunk = np.pad(chunk, (0, chunk_len - len(chunk)))
-                audio_t = torch.tensor(chunk, dtype=torch.float32).unsqueeze(0).to(DEVICE)
-                out = model(enc['input_ids'], enc['attention_mask'], audio_t, text_missing=is_inst)
-                for col in TARGET_COLS: all_logits[col].append(out[col][0])
-        
-        final_results = []
-        encoders = {'mbti': le_mbti, 'emotion': le_emotion, 'vibe': le_vibe, 'intensity': le_intensity, 'tempo': le_tempo}
-        
-        for col in TARGET_COLS:
-            logits_stack = torch.stack(all_logits[col])
-            # Pake mean tanpa weight manual biar model audio kerja murni
-            final_logits = logits_stack.mean(dim=0) / 0.7 
-            
-            if col == 'tempo':
-                t_cls = list(le_tempo.classes_)
-                try:
-                    if tempo_bpm > 125: final_logits[t_cls.index('Fast')] += 3.0
-                except: pass
-
-            probs = torch.nn.functional.softmax(final_logits, dim=0).cpu().numpy()
-            classes = encoders[col].classes_
-            res_dict = {str(classes[i]): float(probs[i]) for i in range(len(classes))}
-            final_results.append(dict(sorted(res_dict.items(), key=lambda x: x[1], reverse=True)[:3]))
+                iv = torch.tensor(chunk).unsqueeze(0).to(DEVICE)
+                aam = torch.ones_like(iv).to(DEVICE)
+                
+                m, e, v, it, t = audio_model(iv, aam)
+                a_mbti.append(F.softmax(m, dim=1))
+                a_emo.append(F.softmax(e, dim=1))
+                a_vibe.append(F.softmax(v, dim=1))
+                a_int.append(F.softmax(it, dim=1))
+                a_tmp.append(t * 200.0) # Denormalize tempo
+
+        avg_a_mbti = torch.stack(a_mbti).mean(dim=0)
+        avg_a_emo = torch.stack(a_emo).mean(dim=0)
+        avg_a_vibe = torch.stack(a_vibe).mean(dim=0)
+        avg_a_int = torch.stack(a_int).mean(dim=0)
+        avg_tempo_bpm = torch.stack(a_tmp).mean().item()
+
+        has_lyrics = lyrics_input and len(str(lyrics_input).strip()) > 15
+        if has_lyrics:
+            t_inputs = tokenizer(str(lyrics_input), truncation=True, padding=True, max_length=256, return_tensors="pt").to(DEVICE)
+            with torch.no_grad():
+                t_mbti_probs = F.softmax(text_mbti_model(**t_inputs).logits, dim=1)
+                t_emo_probs = F.softmax(text_emo_model(**t_inputs).logits, dim=1)
             
-        return final_results
+            final_mbti_probs = (avg_a_mbti * 0.6) + (t_mbti_probs * 0.4)
+            final_emo_probs = (avg_a_emo * 0.6) + (t_emo_probs * 0.4)
+        else:
+            final_mbti_probs = avg_a_mbti
+            final_emo_probs = avg_a_emo
+
+        def process_probs(probs, labels):
+            p = probs.cpu().squeeze().numpy()
+            res = {labels[i]: float(p[i]) for i in range(len(labels))}
+            return dict(sorted(res.items(), key=lambda x: x[1], reverse=True)[:3])
+
+        return [
+            process_probs(final_mbti_probs, MBTI_LABELS),
+            process_probs(final_emo_probs, EMO_LABELS),
+            process_probs(avg_a_vibe, VIBE_LABELS),
+            process_probs(avg_a_int, INTENSITY_LABELS),
+            f"{avg_tempo_bpm:.2f} BPM"
+        ]
+
     except Exception as e: return [{"Error": str(e)}] * 5
 
-# -- INTERFACE --
-with gr.Blocks(theme=gr.themes.Monochrome()) as demo:
-    gr.Markdown("# Neural Math Rock Multimodal Analysis")
-    gr.Markdown("Identify personality and emotional states from music audio and lyrics.")
+# -- GRADIO INTERFACE --
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# 🎸 Neural Math Rock - Hybrid Analysis")
+    gr.Markdown("Identify personality and emotional states using Ensemble Audio-Text Models.")
     
     with gr.Row():
         with gr.Column():
-            search_box = gr.Textbox(label="YouTube Search (Artist - Song Title)", placeholder="Enter song name...")
-            fetch_btn = gr.Button("FETCH AUDIO AND LYRICS", variant="secondary")
-            gr.HTML("<hr>")
-            audio_box = gr.Audio(type="filepath", label="Audio Source")
-            lyrics_box = gr.Textbox(lines=6, label="Lyrics Source", placeholder="Lyrics...")
-            run_btn = gr.Button("RUN ANALYSIS", variant="primary")
-        
+            search_input = gr.Textbox(label="YouTube Search", placeholder="Artist - Song Title")
+            btn_fetch = gr.Button("FETCH ASSETS", variant="secondary")
+            audio_input = gr.Audio(type="filepath", label="Audio Source")
+            lyrics_input = gr.Textbox(lines=6, label="Lyrics Source", placeholder="Paste lyrics for better accuracy...")
+            btn_run = gr.Button("RUN HYBRID ANALYSIS", variant="primary")
+            
         with gr.Column():
-            res_mbti = gr.Label(label="Personality (MBTI)")
-            res_emo = gr.Label(label="Emotional State")
-            res_vibe = gr.Label(label="Acoustic Vibe")
-            res_int = gr.Label(label="Intensity Level")
-            res_tmp = gr.Label(label="Tempo Classification")
-
-    fetch_btn.click(fn=search_and_fetch, inputs=[search_box], outputs=[audio_box, lyrics_box])
-    run_btn.click(fn=analyze_track, inputs=[audio_box, lyrics_box], outputs=[res_mbti, res_emo, res_vibe, res_int, res_tmp])
+            out_mbti = gr.Label(label="Personality (MBTI)")
+            out_emo = gr.Label(label="Emotional State")
+            out_vibe = gr.Label(label="Acoustic Vibe")
+            out_int = gr.Label(label="Intensity Level")
+            out_tmp = gr.Textbox(label="Estimated Tempo (BPM)")
+
+    btn_fetch.click(fn=search_and_fetch, inputs=[search_input], outputs=[audio_input, lyrics_input])
+    btn_run.click(fn=analyze_track, inputs=[audio_input, lyrics_input], outputs=[out_mbti, out_emo, out_vibe, out_int, out_tmp])
 
 if __name__ == "__main__":
     demo.launch()