Update app.py

app.py

@@ -12,12 +12,17 @@ class MultiTaskModel(nn.Module):
         self.topik_classifier = nn.Linear(hidden_size, num_topic_classes)
         self.sentiment_classifier = nn.Linear(hidden_size, num_sentiment_classes)
 
-    def forward(self, input_ids, attention_mask):
-        outputs = self.encoder(input_ids=input_ids, attention_mask=attention_mask)
-        pooled_output = outputs.last_hidden_state[:, 0]
-        topik_logits = self.topik_classifier(pooled_output)
-        sentimen_logits = self.sentimen_classifier(pooled_output)
-        return topik_logits, sentimen_logits
+    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 tokenizer & model
 tokenizer = AutoTokenizer.from_pretrained("tokenizer")