Update app.py

app.py

@@ -3,7 +3,7 @@ import torch.nn as nn
 from transformers import AutoTokenizer, AutoModel
 import gradio as gr
 
-# Model multitask (Topik & Sentimen) dengan nama layer sesuai model.pt
+# 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__()
@@ -13,17 +13,15 @@ class MultiTaskModel(nn.Module):
         self.sentiment_classifier = nn.Linear(hidden_size, num_sentiment_classes)
 
     def forward(self, input_ids, attention_mask, token_type_ids=None):
-    outputs = self.encoder(
-        input_ids=input_ids,
-        attention_mask=attention_mask,
-        token_type_ids=token_type_ids
-    )
-    pooled_output = outputs.last_hidden_state[:, 0]
-    topik_logits = self.topik_classifier(pooled_output)
-    sentimen_logits = self.sentiment_classifier(pooled_output)
-    return topik_logits, sentimen_logits
-
-
+        outputs = self.encoder(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            token_type_ids=token_type_ids
+        )
+        pooled_output = outputs.last_hidden_state[:, 0]
+        topik_logits = self.topik_classifier(pooled_output)
+        sentimen_logits = self.sentiment_classifier(pooled_output)
+        return topik_logits, sentimen_logits
 
 # Load tokenizer & model
 tokenizer = AutoTokenizer.from_pretrained("tokenizer")
@@ -32,7 +30,7 @@ model.load_state_dict(torch.load("model.pt", map_location=torch.device("cpu")))
 model.eval()
 
 # Label mapping
-topik_labels = ["Produk", "Layanan", "Pengiriman", "Lainnya"]
+topik_labels = ["Produk", "Layanan", "Pengiriman", "Pembatalan", "Lainnya"]
 sentimen_labels = ["Negatif", "Netral", "Positif"]
 
 def klasifikasi(text):