Update app.py

app.py

@@ -1,235 +1,214 @@
+import os
+import gc
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-import librosa
+import soundfile as sf
 import torchaudio
 import gradio as gr
+import numpy as np
+from huggingface_hub import hf_hub_download
 from transformers import (
+    AutoModel, 
+    AutoFeatureExtractor, 
+    AutoTokenizer,
     XLMRobertaForSequenceClassification, 
-    XLMRobertaTokenizer, 
-    WavLMModel,
-    AutoFeatureExtractor,
-    ClapAudioModel,
-    ClapProcessor
+    XLMRobertaTokenizer
 )
-from huggingface_hub import hf_hub_download
 import warnings
 
 warnings.filterwarnings('ignore')
 
-# -- CONFIGURATION --
-REPO_AUDIO = "anggars/neural-mathrock"
+# ------------------------------------------------------------------------------
+# CONFIGURATION & LABEL MAPPING
+# ------------------------------------------------------------------------------
+REPO_MAIN      = "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")
+REPO_TEXT_EMO  = "anggars/xlm-emotion"
+MODEL_FILE     = "model.pth"
+DEVICE         = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+CROP_SEC       = 5
+SR_TARGET      = 24000
 
-# -- GLOBAL LABELS --
 MBTI_LABELS = sorted(["INTJ", "INTP", "ENTJ", "ENTP", "INFJ", "INFP", "ENFJ", "ENFP", "ISTJ", "ISFJ", "ESTJ", "ESFJ", "ISTP", "ISFP", "ESTP", "ESFP"])
-EMO_LABELS = ['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 = ['aggressive', 'atmospheric', 'melancholic', 'technical']
-INTENSITY_LABELS = ['low', 'medium', 'high']
-TEMPO_LABELS = ['slow', 'moderate', 'fast']
+EMOTION_CLASSES_AUDIO = ["anger", "disgust", "fear", "joy", "neutral", "sadness", "surprise"]
+EMOTION_CLASSES_TEXT  = ['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']
 
-WAVLM_SR = 16000
-CLAP_SR = 48000
+text_emo2id = {e: i for i, e in enumerate(EMOTION_CLASSES_TEXT)}
 
-emo2id = {e: i for i, e in enumerate(EMO_LABELS)}
+# Calibration Weights derived from SOTA Report 3
+RAW_CE_WEIGHTS = [0.563, 16.326, 4.525, 0.968, 2.704, 0.473, 0.702]
+SOFT_WEIGHTS   = torch.sqrt(torch.tensor(RAW_CE_WEIGHTS, dtype=torch.float32)).to(DEVICE)
+TEMPERATURE    = 2.0
 
-# -- HYBRID ARCHITECTURE ALIGNED WITH TRAINING EXECUTION --
-class AudioMathRockModel(nn.Module):
-    def __init__(self):
+# ------------------------------------------------------------------------------
+# NEURAL ARCHITECTURE (7-CLASS AUDIO FUSION)
+# ------------------------------------------------------------------------------
+class MultimodalFusionClassifier(nn.Module):
+    def __init__(self, num_classes=7):
         super().__init__()
-        self.wavlm = WavLMModel.from_pretrained("microsoft/wavlm-base")
-        self.clap  = ClapAudioModel.from_pretrained("laion/clap-htsat-unfused")
-        
-        # 2D-CNN Mel Spectrogram Extractor Block
-        self.mel_spectrogram = torchaudio.transforms.MelSpectrogram(sample_rate=WAVLM_SR, n_fft=1024, hop_length=512, n_mels=128)
-        self.amplitude_to_db = torchaudio.transforms.AmplitudeToDB(stype='power', top_db=80.0)
-        
-        self.cnn_extractor = nn.Sequential(
-            nn.Conv2d(1, 32, kernel_size=3, stride=1, padding=1),
-            nn.BatchNorm2d(32), nn.ReLU(), nn.MaxPool2d(2, 2),
-            nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1),
-            nn.BatchNorm2d(64), nn.ReLU(), nn.MaxPool2d(2, 2),
-            nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1),
-            nn.BatchNorm2d(128), nn.ReLU(), nn.AdaptiveAvgPool2d((1, 1)),
-            nn.Flatten(),
-            nn.Linear(128, 512)
+        self.audio_model = AutoModel.from_pretrained("m-a-p/MERT-v1-330M", trust_remote_code=True)
+        self.text_model = AutoModel.from_pretrained("FacebookAI/roberta-base")
+        self.fusion_head = nn.Sequential(
+            nn.Linear(1024 + 768, 512),
+            nn.LayerNorm(512),
+            nn.GELU(),
+            nn.Dropout(0.5),
+            nn.Linear(512, 256),
+            nn.LayerNorm(256),
+            nn.GELU(),
+            nn.Dropout(0.4),
+            nn.Linear(256, num_classes)
         )
-        
-        self.wlm_proj = nn.Linear(768, 512)
-        self.clp_proj = nn.Linear(768, 512)
-        
-        # Expanded dimension to 1536 to hold CNN features
-        self.fusion = nn.Sequential(
-            nn.Linear(1536, 512), 
-            nn.LayerNorm(512), 
-            nn.Tanh(), 
-            nn.Dropout(0.3)
-        )
-        
-        self.emo_head = nn.Sequential(
-            nn.Linear(512, 256), nn.GELU(), nn.Dropout(0.2), 
-            nn.Linear(256, len(EMO_LABELS))
-        )
-        self.vibe_head = nn.Linear(512, len(VIBE_LABELS))
-        self.int_head  = nn.Linear(512, len(INTENSITY_LABELS))
-        self.tmp_head  = nn.Linear(512, len(TEMPO_LABELS))
 
-    def forward(self, wavlm_values, clap_values):
-        wavlm_feats = self.wavlm(wavlm_values).last_hidden_state.mean(dim=1)
-        clap_feats  = self.clap(clap_values).pooler_output
-        
-        with torch.no_grad():
-            mel = self.mel_spectrogram(wavlm_values.float())
-            mel_db = self.amplitude_to_db(mel)
-        
-        cnn_feats = F.gelu(self.cnn_extractor(mel_db.unsqueeze(1))).to(wavlm_feats.dtype)
-        wlm_p = F.gelu(self.wlm_proj(wavlm_feats))
-        clp_p = F.gelu(self.clp_proj(clap_feats))
+    def forward(self, audio_vals, text_ids, text_mask):
+        audio_feats = self.audio_model(audio_vals, output_hidden_states=True).last_hidden_state.mean(dim=1)
+        text_feats = self.text_model(input_ids=text_ids, attention_mask=text_mask).last_hidden_state.mean(dim=1)
+        return self.fusion_head(torch.cat((audio_feats, text_feats), dim=-1))
+
+# ------------------------------------------------------------------------------
+# MODEL INITIALIZATION
+# ------------------------------------------------------------------------------
+print("[SYSTEM] Booting Models into Memory...")
+mert_ext = AutoFeatureExtractor.from_pretrained("m-a-p/MERT-v1-330M", trust_remote_code=True)
+roberta_tok = AutoTokenizer.from_pretrained("FacebookAI/roberta-base")
+
+fusion_model = MultimodalFusionClassifier(num_classes=len(EMOTION_CLASSES_AUDIO)).to(DEVICE)
+ckpt_path = hf_hub_download(repo_id=REPO_MAIN, filename=MODEL_FILE)
+fusion_model.load_state_dict(torch.load(ckpt_path, map_location=DEVICE)['model_state_dict'])
+fusion_model.eval()
+
+# Load Text-Only Models (Isolated)
+mbti_tok = XLMRobertaTokenizer.from_pretrained(REPO_TEXT_MBTI)
+mbti_model = XLMRobertaForSequenceClassification.from_pretrained(REPO_TEXT_MBTI).to(DEVICE)
+mbti_model.eval()
+
+# Load High-Resolution 28-Class Text Emotion Model
+emo28_model = XLMRobertaForSequenceClassification.from_pretrained(REPO_TEXT_EMO).to(DEVICE)
+emo28_model.eval()
+
+# ------------------------------------------------------------------------------
+# INFERENCE ENGINE
+# ------------------------------------------------------------------------------
+def process_audio_chunks(audio_path):
+    waveform_np, orig_sr = sf.read(audio_path)
+    if len(waveform_np.shape) > 1: 
+        waveform_np = waveform_np.mean(axis=1) 
+    
+    waveform = torch.tensor(waveform_np, dtype=torch.float32).unsqueeze(0)
+    if orig_sr != SR_TARGET:
+        waveform = torchaudio.functional.resample(waveform, orig_sr, SR_TARGET)
+    
+    waveform = waveform.squeeze(0)
+    chunk_samples = SR_TARGET * CROP_SEC
+    
+    chunks = [waveform[i:i+chunk_samples].numpy() for i in range(0, len(waveform), chunk_samples) if len(waveform[i:i+chunk_samples]) > SR_TARGET * 2]
+    if not chunks:
+        chunks = [F.pad(waveform, (0, chunk_samples - waveform.shape[0])).numpy()]
         
-        fused = self.fusion(torch.cat([wlm_p, clp_p, cnn_feats], dim=-1))
-        return self.emo_head(fused), self.vibe_head(fused), self.int_head(fused), self.tmp_head(fused)
-
-# -- MODEL INITIALIZATION --
-print("Fetching Model Weights...")
-ckpt_path = hf_hub_download(repo_id=REPO_AUDIO, filename="model.pth")
-ckpt = torch.load(ckpt_path, map_location=DEVICE, weights_only=False)
+    return chunks[:10]
 
-audio_model = AudioMathRockModel().to(DEVICE)
-audio_model.load_state_dict(ckpt['model_state_dict'], strict=True)
-audio_model.eval()
-
-wavlm_extractor = AutoFeatureExtractor.from_pretrained("microsoft/wavlm-base")
-clap_processor  = ClapProcessor.from_pretrained("laion/clap-htsat-unfused")
-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()
-
-# -- ANALYSIS ENGINE --
 def analyze_track(audio_path, lyrics_input):
     has_audio = audio_path is not None
-    has_lyrics = lyrics_input is not None and len(str(lyrics_input).strip()) > 15
-
+    safe_lyrics = str(lyrics_input).strip() if lyrics_input else ""
+    has_lyrics = len(safe_lyrics) > 10
+    
     if not has_audio and not has_lyrics:
-        return {"Error": 1.0}, {"Error": 1.0}, {"Error": 1.0}, {"Error": 1.0}, {"Error": 1.0}
+        return {"Error": 1.0}, {"Error": 1.0}
 
-    res_mbti, res_emo, res_vibe, res_int, res_tmp = {}, {}, {}, {}, {}
+    res_mbti = {}
+    res_emo_final = {}
 
-    # --- TEXT ONLY PROCESSING ---
-    if has_lyrics and not has_audio:
-        t_in = tokenizer(str(lyrics_input), truncation=True, padding=True, max_length=256, return_tensors="pt").to(DEVICE)
+    # 1. ISOLATED MBTI & 28-CLASS TEXT INFERENCE
+    t_emo28_probs = None
+    if has_lyrics:
+        t_in = mbti_tok(safe_lyrics, 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_in).logits, dim=1).cpu().squeeze().numpy()
-            t_emo_probs = F.softmax(text_emo_model(**t_in).logits, dim=1).cpu().squeeze().numpy()
-        
-        mbti_dict = {MBTI_LABELS[i]: float(t_mbti_probs[i]) for i in range(len(MBTI_LABELS))}
-        emo_dict = {EMO_LABELS[i]: float(t_emo_probs[i]) for i in range(len(EMO_LABELS))}
-        
+            mbti_probs = F.softmax(mbti_model(**t_in).logits, dim=1).cpu().squeeze().numpy()
+            t_emo28_probs = F.softmax(emo28_model(**t_in).logits, dim=1).cpu().squeeze().numpy()
+            
+        mbti_dict = {MBTI_LABELS[i]: float(mbti_probs[i]) for i in range(len(MBTI_LABELS))}
         res_mbti = dict(sorted(mbti_dict.items(), key=lambda x: x[1], reverse=True)[:3])
-        res_emo = dict(sorted(emo_dict.items(), key=lambda x: x[1], reverse=True)[:3])
-        return res_mbti, res_emo, {}, {}, {}
-
-    # --- AUDIO PROCESSING ---
-    try:
-        wav_orig, orig_sr = librosa.load(audio_path, sr=None, mono=True)
-        waveform = torch.tensor(wav_orig).unsqueeze(0)
-        
-        wlm_wave = torchaudio.functional.resample(waveform, orig_sr, WAVLM_SR).squeeze(0)
-        clp_wave = torchaudio.functional.resample(waveform, orig_sr, CLAP_SR).squeeze(0)
-
-        wavlm_samples = WAVLM_SR * 15
-        clap_samples = CLAP_SR * 15
+    else:
+        res_mbti = {"Requires lyrics for MBTI": 1.0}
 
-        wlm_chunks = [wlm_wave[i:i+wavlm_samples].numpy() for i in range(0, len(wlm_wave), wavlm_samples) if len(wlm_wave[i:i+wavlm_samples]) > WAVLM_SR]
-        clp_chunks = [clp_wave[i:i+clap_samples].numpy() for i in range(0, len(clp_wave), clap_samples) if len(clp_wave[i:i+clap_samples]) > CLAP_SR]
-        
-        if not wlm_chunks:
-            wlm_chunks = [F.pad(wlm_wave, (0, wavlm_samples - wlm_wave.shape[0])).numpy()]
-            clp_chunks = [F.pad(clp_wave, (0, clap_samples - clp_wave.shape[0])).numpy()]
-
-        el_list, vl_list, il_list, tl_list = [], [], [], []
-        
-        with torch.no_grad():
-            for w_c, c_c in zip(wlm_chunks, clp_chunks):
-                wlm_inputs = wavlm_extractor([w_c], sampling_rate=WAVLM_SR, return_tensors="pt")["input_values"].to(DEVICE)
-                clp_inputs = clap_processor(audio=[c_c], sampling_rate=CLAP_SR, return_tensors="pt")["input_features"].to(DEVICE)
-                
-                el, vl, il, tl = audio_model(wlm_inputs, clp_inputs)
-                el_list.append(el.cpu())
-                vl_list.append(vl.cpu())
-                il_list.append(il.cpu())
-                tl_list.append(tl.cpu())
-
-        raw_el = torch.cat(el_list, dim=0).mean(dim=0, keepdim=True)
-        raw_vl = torch.cat(vl_list, dim=0).mean(dim=0, keepdim=True)
-        raw_il = torch.cat(il_list, dim=0).mean(dim=0, keepdim=True)
-        raw_tl = torch.cat(tl_list, dim=0).mean(dim=0, keepdim=True)
-
-        # Enforce highly competitive soft-scaling for audio emotion logits
-        scaled_el = raw_el / raw_el.std(dim=-1, keepdim=True).clamp(min=1e-6)
-        probs_el_audio = F.softmax(scaled_el / 0.3, dim=-1).squeeze().numpy()
-        
-        probs_vl = F.softmax(raw_vl, dim=-1).squeeze().numpy()
-        probs_il = F.softmax(raw_il, dim=-1).squeeze().numpy()
-        probs_tl = F.softmax(raw_tl, dim=-1).squeeze().numpy()
-
-        res_emo_raw = {EMO_LABELS[i]: float(probs_el_audio[i]) for i in range(len(EMO_LABELS))}
-        res_vibe = {VIBE_LABELS[i]: float(probs_vl[i]) for i in range(len(VIBE_LABELS))}
-        res_int = {INTENSITY_LABELS[i]: float(probs_il[i]) for i in range(len(INTENSITY_LABELS))}
-        res_tmp = {TEMPO_LABELS[i]: float(probs_tl[i]) for i in range(len(TEMPO_LABELS))}
-
-        res_vibe = dict(sorted(res_vibe.items(), key=lambda x: x[1], reverse=True)[:3])
-        res_int = dict(sorted(res_int.items(), key=lambda x: x[1], reverse=True)[:3])
-        res_tmp = dict(sorted(res_tmp.items(), key=lambda x: x[1], reverse=True)[:3])
-
-        # --- MULTIMODAL LATE FUSION ENGAGEMENT ---
-        if has_lyrics:
-            t_in = tokenizer(str(lyrics_input), truncation=True, padding=True, max_length=256, return_tensors="pt").to(DEVICE)
+    # 2. AUDIO FUSION INFERENCE
+    if has_audio:
+        try:
+            audio_chunks = process_audio_chunks(audio_path)
+            
+            t_inputs = roberta_tok([safe_lyrics] * len(audio_chunks), padding=True, truncation=True, max_length=128, return_tensors="pt")
+            t_ids = t_inputs["input_ids"].to(DEVICE)
+            t_mask = t_inputs["attention_mask"].to(DEVICE)
+            
+            a_inputs = mert_ext(audio_chunks, sampling_rate=SR_TARGET, return_tensors="pt", padding="max_length", truncation=True, max_length=SR_TARGET * CROP_SEC)
+            a_vals = a_inputs["input_values"].to(DEVICE)
+            
             with torch.no_grad():
-                t_mbti_probs = F.softmax(text_mbti_model(**t_in).logits, dim=1).cpu().squeeze().numpy()
-                t_emo_probs = F.softmax(text_emo_model(**t_in).logits, dim=1).cpu().squeeze().numpy()
+                logits = fusion_model(a_vals, t_ids, t_mask)
+                avg_logits = logits.mean(dim=0)
+                
+                scaled_logits = avg_logits / TEMPERATURE
+                probs = F.softmax(scaled_logits, dim=-1)
+                adjusted_probs = probs * SOFT_WEIGHTS
+                audio_7_probs = (adjusted_probs / adjusted_probs.sum()).cpu().numpy()
+                
+            audio_emo_dict = {EMOTION_CLASSES_AUDIO[i]: float(audio_7_probs[i]) for i in range(len(EMOTION_CLASSES_AUDIO))}
             
-            m_dict = {MBTI_LABELS[i]: float(t_mbti_probs[i]) for i in range(len(MBTI_LABELS))}
-            res_mbti = dict(sorted(m_dict.items(), key=lambda x: x[1], reverse=True)[:3])
+            # --- TRUE MULTIMODAL LATE FUSION ENGAGEMENT ---
+            if has_lyrics:
+                # Bikin dictionary kosong untuk 28 emosi
+                fusion_28_dict = {label: 0.0 for label in EMOTION_CLASSES_TEXT}
+                
+                # Masukin probabilitas murni dari Teks (Bobot 60%)
+                for i, label in enumerate(EMOTION_CLASSES_TEXT):
+                    fusion_28_dict[label] += float(t_emo28_probs[i]) * 0.60
+                    
+                # Suntikin probabilitas dari Audio 7 Kelas ke 28 Kelas Teks (Bobot 40%)
+                for label in EMOTION_CLASSES_AUDIO:
+                    if label in fusion_28_dict:
+                        fusion_28_dict[label] += audio_emo_dict[label] * 0.40
+                        
+                # Normalisasi ulang biar totalnya 1.0
+                total_prob = sum(fusion_28_dict.values())
+                fusion_28_dict = {k: v / total_prob for k, v in fusion_28_dict.items()}
+                
+                res_emo_final = dict(sorted(fusion_28_dict.items(), key=lambda x: x[1], reverse=True)[:5])
+            else:
+                # Unimodal Audio (Hanya ngeluarin 7 Kelas)
+                res_emo_final = dict(sorted(audio_emo_dict.items(), key=lambda x: x[1], reverse=True)[:4])
             
-            # Weighted late fusion to prioritize text model for emotion taxonomy resolution
-            for k in res_emo_raw.keys():
-                idx = emo2id[k]
-                res_emo_raw[k] = (res_emo_raw[k] * 0.15) + (float(t_emo_probs[idx]) * 0.85)
-            res_emo = dict(sorted(res_emo_raw.items(), key=lambda x: x[1], reverse=True)[:3])
-        else:
-            res_mbti = {}
-            res_emo = dict(sorted(res_emo_raw.items(), key=lambda x: x[1], reverse=True)[:3])
-
-        return res_mbti, res_emo, res_vibe, res_int, res_tmp
-
-    except Exception as e:
-        print(f"PIPELINE ERROR: {str(e)}") 
-        return {"System Error": 1.0}, {"System Error": 1.0}, {"System Error": 1.0}, {"System Error": 1.0}, {"System Error": 1.0}
-
-# -- INTERFACE BLOCK --
+        except Exception as e:
+            res_emo_final = {f"Audio Error: {str(e)}": 1.0}
+    
+    # 3. TEXT-ONLY FALLBACK (Kalau audio gak dimasukin)
+    elif has_lyrics:
+        emo_dict = {EMOTION_CLASSES_TEXT[i]: float(t_emo28_probs[i]) for i in range(len(EMOTION_CLASSES_TEXT))}
+        res_emo_final = dict(sorted(emo_dict.items(), key=lambda x: x[1], reverse=True)[:5])
+
+    return res_mbti, res_emo_final
+
+# ------------------------------------------------------------------------------
+# GRADIO 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 pure music audio and lyrics inferences.")
+    gr.Markdown("Identify personality (Text) and emotional states (MERT+RoBERTa Multimodal) from Math Rock & Midwest Emo tracks.")
     
     with gr.Row():
         with gr.Column():
-            audio_box = gr.Audio(type="filepath", label="Audio Source")
-            lyrics_box = gr.Textbox(lines=8, label="Lyrics Source", placeholder="Paste lyrics here for hybrid analysis...")
-            run_btn = gr.Button("RUN ANALYSIS", variant="primary")
+            audio_box = gr.Audio(type="filepath", label="Audio Source (.wav / .mp3)")
+            lyrics_box = gr.Textbox(lines=8, label="Lyrics Source", placeholder="Paste lyrics here... \n(If blank: Analyzes 7 Audio Emotions. If filled: Analyzes 28 High-Resolution Emotions.)")
+            run_btn = gr.Button("RUN SOTA 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")
+            res_mbti = gr.Label(label="Personality (MBTI - Text Only)")
+            res_emo = gr.Label(label="Emotional State (Multimodal 28-Class Fusion)")
 
     run_btn.click(
         fn=analyze_track, 
         inputs=[audio_box, lyrics_box], 
-        outputs=[res_mbti, res_emo, res_vibe, res_int, res_tmp]
+        outputs=[res_mbti, res_emo]
     )
 
 if __name__ == "__main__":