Update app.py

app.py

@@ -3,56 +3,55 @@ import torch.nn as nn
 from transformers import AutoTokenizer, AutoModel, AutoModelForSeq2SeqLM
 import gradio as gr
 
-# Model klasifikasi multitugas
+# Definisi model klasifikasi multitugas
 class MultiTaskModel(nn.Module):
     def __init__(self, base_model_name, num_topic_classes, num_sentiment_classes):
-        super(MultiTaskModel, self).__init__()
+        super().__init__()
         self.encoder = AutoModel.from_pretrained(base_model_name)
-        hidden_size = self.encoder.config.hidden_size
-        self.topik_classifier = nn.Linear(hidden_size, num_topic_classes)
-        self.sentiment_classifier = nn.Linear(hidden_size, num_sentiment_classes)
+        hs = self.encoder.config.hidden_size
+        self.topik_classifier = nn.Linear(hs, num_topic_classes)
+        self.sentiment_classifier = nn.Linear(hs, 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
-
-# Load model klasifikasi
-tokenizer = AutoTokenizer.from_pretrained("tokenizer")
+        out = self.encoder(input_ids=input_ids,
+                           attention_mask=attention_mask,
+                           token_type_ids=token_type_ids)
+        pooled = out.last_hidden_state[:, 0]
+        return self.topik_classifier(pooled), self.sentiment_classifier(pooled)
+
+# Load tokenizer dan model klasifikasi
+tokenizer = AutoTokenizer.from_pretrained("tokenizer")  # Folder tokenizer harus diupload
 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.load_state_dict(torch.load("model.pt", map_location="cpu"))
 model.eval()
 
-# Load model summarization
-sum_tokenizer = AutoTokenizer.from_pretrained("cahya/bart-base-indonesian-summarization")
-sum_model = AutoModelForSeq2SeqLM.from_pretrained("cahya/bart-base-indonesian-summarization")
+# Load tokenizer dan model summarization
+sum_tok = AutoTokenizer.from_pretrained("xTorch8/bart-id-summarization")
+sum_model = AutoModelForSeq2SeqLM.from_pretrained("xTorch8/bart-id-summarization")
 
-topik_labels = ["Produk", "Layanan", "Pengiriman", "Pembatalan", "Lainnya"]
-sentimen_labels = ["Negatif", "Netral", "Positif"]
+# Label klasifikasi
+labels_topik = ["Produk", "Layanan", "Pengiriman", "Pembatalan", "Lainnya"]
+labels_sentiment = ["Negatif", "Netral", "Positif"]
 
-def klasifikasi(text):
-    # Klasifikasi
-    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
+# Fungsi analisis
+def analyze(text):
+    # Klasifikasi topik & sentimen
+    inp = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
     with torch.no_grad():
-        topik_logits, sentimen_logits = model(**inputs)
-        topik_idx = torch.argmax(topik_logits, dim=-1).item()
-        sentimen_idx = torch.argmax(sentimen_logits, dim=-1).item()
-
-    topik = topik_labels[topik_idx]
-    sentimen = sentimen_labels[sentimen_idx]
-
-    # Ringkasan
-    sum_inputs = sum_tokenizer(text, return_tensors="pt", max_length=512, truncation=True)
-    summary_ids = sum_model.generate(**sum_inputs, max_length=40, min_length=10, do_sample=False)
-    ringkasan = sum_tokenizer.decode(summary_ids[0], skip_special_tokens=True)
-
-    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")
+        t_logits, s_logits = model(**inp)
+        topik = labels_topik[int(torch.argmax(t_logits))]
+        sentimen = labels_sentiment[int(torch.argmax(s_logits))]
+
+    # Ringkasan teks
+    s_inp = sum_tok(text, return_tensors="pt", truncation=True, padding=True)
+    summ_ids = sum_model.generate(**s_inp, max_length=50, num_beams=2)
+    ringkasan = sum_tok.decode(summ_ids[0], skip_special_tokens=True)
+
+    return (f"HASIL ANALISIS\n"
+            f"Topik: {topik}\n"
+            f"Sentimen: {sentimen}\n"
+            f"Ringkasan: {ringkasan}")
+
+# Gradio interface
+demo = gr.Interface(fn=analyze, inputs="text", outputs="text", title="Analisis Topik, Sentimen, dan Ringkasan Pelanggan")
 demo.launch()