app.py

import gradio as gr
from transformers import pipeline
from deep_translator import GoogleTranslator
from langdetect import detect, DetectorFactory
import math

DetectorFactory.seed = 0

print("Loading Models...")
# Pake top_k=None biar semua label dapet score buat kalkulasi rata-rata chunk
emotion_pipe = pipeline("text-classification", model="anggars/xlm-emotion", top_k=None)
mbti_pipe = pipeline("text-classification", model="anggars/xlm-mbti", top_k=None)

MANUAL_MAPPING = {
    "sombong": "DISAPPROVAL", "arogan": "DISAPPROVAL", "belagu": "DISAPPROVAL",
    "geuleuh": "DISGUST", "hareudang": "NERVOUSNESS", "gandeng": "ANNOYANCE",
    "nuhun": "GRATITUDE", "hampura": "REMORSE", "bungah": "JOY",
    "keuheul": "ANGER", "goblog": "ANGER", "anjrit": "SURPRISE",
    "ngelu": "SADNESS", "mangkel": "ANGER", "guyon": "AMUSEMENT",
    "gabut": "SADNESS", "bacot": "ANNOYANCE"
}

def chunk_text(text, max_words=300):
    """Mecah teks panjang jadi potongan-potongan biar gak error 512 token"""
    words = text.split()
    chunks = []
    for i in range(0, len(words), max_words):
        chunk = " ".join(words[i:i+max_words])
        chunks.append(chunk)
    return chunks

def predict_logic(text):
    if not text or text.strip() == "": return {}, {}
    original_text = text.lower()
    
    try:
        lang = detect(text)
    except:
        lang = "id"

    # Terjemahkan kalau bukan Indo/Inggris
    text_ready = GoogleTranslator(source='auto', target='en').translate(text) if lang not in ['id', 'en'] else text

    chunks = chunk_text(text_ready, max_words=250) 
    
    temp_emo_scores = {}
    temp_mbti_scores = {}
    
    for chunk in chunks:
        # Inference pake model yang udah lu benerin
        emo_res = emotion_pipe(chunk, truncation=True, max_length=512)[0]
        mbti_res = mbti_pipe(chunk, truncation=True, max_length=512)[0]
        
        for item in emo_res:
            temp_emo_scores[item['label']] = temp_emo_scores.get(item['label'], 0) + item['score']
        for item in mbti_res:
            temp_mbti_scores[item['label']] = temp_mbti_scores.get(item['label'], 0) + item['score']

    num_chunks = len(chunks)
    final_emo = {k: v / num_chunks for k, v in temp_emo_scores.items()}
    final_mbti = {k: v / num_chunks for k, v in temp_mbti_scores.items()}

    # --- LOGIKA PENALTI MANUAL DICABUT KARENA MODEL SUDAH BALANCED ---
    # Model lu sekarang udah akurat, gak perlu lagi dikali 0.65 atau 0.55.

    # Manual Mapping buat slang tetep disisain buat bantu deteksi kata lokal
    for word, target_emotion in MANUAL_MAPPING.items():
        if word in original_text:
            if target_emotion in final_emo:
                final_emo[target_emotion] += 0.5 
    
    # Normalisasi score
    total_raw_emo = sum(final_emo.values())
    final_emo = {k: v / total_raw_emo for k, v in final_emo.items()}
    
    total_raw_mbti = sum(final_mbti.values())
    final_mbti = {k: v / total_raw_mbti for k, v in final_mbti.items()}

    # Sorting hasil buat tampilan UI
    sorted_emo = dict(sorted(final_emo.items(), key=lambda item: item[1], reverse=True)[:5])
    sorted_mbti = dict(sorted(final_mbti.items(), key=lambda item: item[1], reverse=True)[:3])

    return sorted_emo, sorted_mbti

# --- UI GRADIO (Copywriting Dijaga Sesuai Request) ---
with gr.Blocks(title="MBTI Emotion Analyzer") as demo:
    gr.Markdown("# MBTI Emotion Intelligence Engine")
    with gr.Row():
        input_txt = gr.Textbox(label="Input Text", placeholder="Type here...", lines=5)
    btn = gr.Button("Analyze Now", variant="primary")
    with gr.Row():
        out_emo = gr.Label(label="Emotion Score")
        out_mbti = gr.Label(label="MBTI Score")
    btn.click(fn=predict_logic, inputs=input_txt, outputs=[out_emo, out_mbti], api_name="predict")

if __name__ == "__main__":
    demo.launch()