Update app.py

app.py

@@ -3,7 +3,6 @@ import torch.nn as nn
 from transformers import AutoTokenizer, AutoModel
 import gradio as gr
 
-# Model multitask dengan token_type_ids support
 class MultiTaskModel(nn.Module):
     def __init__(self, base_model_name, num_topic_classes, num_sentiment_classes):
         super(MultiTaskModel, self).__init__()
@@ -23,13 +22,11 @@ class MultiTaskModel(nn.Module):
         sentimen_logits = self.sentiment_classifier(pooled_output)
         return topik_logits, sentimen_logits
 
-# Load tokenizer & model
 tokenizer = AutoTokenizer.from_pretrained("tokenizer")
 model = MultiTaskModel("indobenchmark/indobert-base-p1", num_topic_classes=5, num_sentiment_classes=3)
 model.load_state_dict(torch.load("model.pt", map_location=torch.device("cpu")))
 model.eval()
 
-# Label mapping
 topik_labels = ["Produk", "Layanan", "Pengiriman", "Pembatalan", "Lainnya"]
 sentimen_labels = ["Negatif", "Netral", "Positif"]
 
@@ -37,12 +34,14 @@ def klasifikasi(text):
     inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
     with torch.no_grad():
         topik_logits, sentimen_logits = model(**inputs)
-        topik_probs = torch.softmax(topik_logits, dim=-1).squeeze()
-        sentimen_probs = torch.softmax(sentimen_logits, dim=-1).squeeze()
+        topik_idx = torch.argmax(topik_logits, dim=-1).item()
+        sentimen_idx = torch.argmax(sentimen_logits, dim=-1).item()
 
-    topik_result = {label: float(prob) for label, prob in zip(topik_labels, topik_probs)}
-    sentimen_result = {label: float(prob) for label, prob in zip(sentimen_labels, sentimen_probs)}
-    return {"Topik": topik_result, "Sentimen": sentimen_result}
+    topik = topik_labels[topik_idx]
+    sentimen = sentimen_labels[sentimen_idx]
+    ringkasan = f"Pelanggan: {text}\nCS: Kami senang dapat membantu Anda."
 
-demo = gr.Interface(fn=klasifikasi, inputs="text", outputs="json", title="Klasifikasi Topik dan Sentimen Pelanggan")
+    return f"HASIL ANALISIS\nTopik: {topik}\nSentimen: {sentimen}\nRingkasan: {ringkasan}"
+
+demo = gr.Interface(fn=klasifikasi, inputs="text", outputs="text", title="Klasifikasi Topik dan Sentimen Pelanggan")
 demo.launch()