Update app.py

app.py

@@ -1,8 +1,9 @@
 import torch
 import torch.nn as nn
-from transformers import AutoTokenizer, AutoModel
+from transformers import AutoTokenizer, AutoModel, AutoModelForSeq2SeqLM
 import gradio as gr
 
+# Model klasifikasi multitugas
 class MultiTaskModel(nn.Module):
     def __init__(self, base_model_name, num_topic_classes, num_sentiment_classes):
         super(MultiTaskModel, self).__init__()
@@ -22,15 +23,21 @@ class MultiTaskModel(nn.Module):
         sentimen_logits = self.sentiment_classifier(pooled_output)
         return topik_logits, sentimen_logits
 
+# Load model klasifikasi
 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()
 
+# Load model summarization
+sum_tokenizer = AutoTokenizer.from_pretrained("cahya/bart-base-indonesian-summarization")
+sum_model = AutoModelForSeq2SeqLM.from_pretrained("cahya/bart-base-indonesian-summarization")
+
 topik_labels = ["Produk", "Layanan", "Pengiriman", "Pembatalan", "Lainnya"]
 sentimen_labels = ["Negatif", "Netral", "Positif"]
 
 def klasifikasi(text):
+    # Klasifikasi
     inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
     with torch.no_grad():
         topik_logits, sentimen_logits = model(**inputs)
@@ -39,7 +46,11 @@ def klasifikasi(text):
 
     topik = topik_labels[topik_idx]
     sentimen = sentimen_labels[sentimen_idx]
-    ringkasan = f"Pelanggan: {text}\nCS: Kami senang dapat membantu Anda."
+
+    # 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}"