8

app.py

@@ -1,6 +1,6 @@
 import gradio as gr
 import torch
-from transformers import BertTokenizer
+from transformers import BertTokenizer, BertModel  # Add BertModel to the imports
 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)
+        self.bert = BertModel.from_pretrained(bert_model_name)  # Use BertModel correctly
         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)