back

app.py

@@ -1,6 +1,6 @@
 import gradio as gr
 import torch
-from transformers import BertTokenizer, BertModel  # Add BertModel to the imports
+from transformers import BertTokenizer, BertModel  
 import torch.nn.functional as F
 
 # Load model and tokenizer from Hugging Face
@@ -10,7 +10,7 @@ tokenizer = BertTokenizer.from_pretrained(model_name)
 class IndoBERTMultiTaskClassifier(torch.nn.Module):
     def __init__(self, bert_model_name, num_labels_task1, num_labels_task2, dropout_rate=0.3):
         super(IndoBERTMultiTaskClassifier, self).__init__()
-        self.bert = BertModel.from_pretrained(bert_model_name)  # Use BertModel correctly
+        self.bert = BertModel.from_pretrained(bert_model_name)
         self.dropout = torch.nn.Dropout(dropout_rate)
         self.classifier_task1 = torch.nn.Linear(self.bert.config.hidden_size, num_labels_task1)
         self.classifier_task2 = torch.nn.Linear(self.bert.config.hidden_size, num_labels_task2)
@@ -25,13 +25,12 @@ class IndoBERTMultiTaskClassifier(torch.nn.Module):
 
         return logits_task1, logits_task2
 
-# Load the model checkpoint into your multitask model class
+# Load model directly from Hugging Face
 model = IndoBERTMultiTaskClassifier(
     bert_model_name=model_name,
     num_labels_task1=3,  # Adjust with your task1 classes
     num_labels_task2=3   # Adjust with your task2 classes
 )
-model.load_state_dict(torch.load("pytorch_model.bin", map_location=torch.device('cpu')))
 model.eval()
 
 # Define label mappings
@@ -52,27 +51,26 @@ def classify(text):
     probs_task1 = F.softmax(logits_task1, dim=1).cpu().numpy()[0]  # Extract the first batch item
     probs_task2 = F.softmax(logits_task2, dim=1).cpu().numpy()[0]  # Extract the first batch item
 
-    # Map probabilities to their corresponding labels
-    result_task1 = {label: prob for label, prob in zip(label_mapping_task1, probs_task1)}
-    result_task2 = {label: prob for label, prob in zip(label_mapping_task2, probs_task2)}
+    # Predict label with highest probability
+    pred_task1 = label_mapping_task1[probs_task1.argmax()]
+    pred_task2 = label_mapping_task2[probs_task2.argmax()]
+
+    # Format probabilities as percentages
+    probs_task1_str = ", ".join([f"{label}: {prob*100:.2f}%" for label, prob in zip(label_mapping_task1, probs_task1)])
+    probs_task2_str = ", ".join([f"{label}: {prob*100:.2f}%" for label, prob in zip(label_mapping_task2, probs_task2)])
+
+    # Combine label predictions with their probabilities
+    result_task1 = f"{pred_task1} ({probs_task1_str})"
+    result_task2 = f"{pred_task2} ({probs_task2_str})"
 
     return result_task1, result_task2
 
-# Gradio Interface with percentage bars
-iface = gr.Interface(
-    fn=classify, 
-    inputs="text", 
-    outputs=[
-        gr.Label(label="Fake Review Detection"), 
-        gr.Label(label="Sentiment Classification")
-    ],
-    title="Multitask IndoBERT: Fake Review & Sentiment Classification",
-    description="Enter a skincare product review in Indonesian and the model will classify it as fake or trusted, and determine the sentiment.",
-    examples=[
-        ["Jokowi sangat kecewa dengan POLRI atas kerusuhan yang terjadi di Malang"],
-        ["Lesti marah terhadap perlakuan KDRT yang dilakukan oleh Bilar"],
-        ["Ungkapan rasa bahagia diutarakan oleh Coki Pardede karena kebebasannya dari penjara"]
-    ]
-)
+# Gradio Interface
+iface = gr.Interface(fn=classify, 
+                     inputs="text", 
+                     outputs=[gr.Label(label="Fake Review Detection"), 
+                              gr.Label(label="Sentiment Classification")],
+                     title="Multitask IndoBERT: Fake Review & Sentiment Classification",
+                     description="Enter a skincare product review in Indonesian and the model will classify it as fake or trusted, and determine the sentiment.")
 
-iface.launch()
+iface.launch()