Update app.py

app.py

@@ -3,7 +3,7 @@ import pandas as pd
 import torch
 from torch import nn
 from torch.utils.data import Dataset, DataLoader
-from transformers import AutoTokenizer, BertModel, BertForSequenceClassification
+from transformers import AutoTokenizer, BertModel
 from sklearn import metrics
 import streamlit as st
 
@@ -90,16 +90,14 @@ class BertClass(torch.nn.Module):
 class PretrainedBertClass(torch.nn.Module):
     def __init__(self):
         super(PretrainedBertClass, self).__init__()
-        self.l1 = BertForSequenceClassification.from_pretrained(bert_path, num_labels=6)
-
-    # return_dict must equal False for Huggingface Transformers v4+
-    def forward(self, input_ids, attention_mask, token_type_ids):
-        _, output = self.l1(
-            input_ids=input_ids,
-            attention_mask=attention_mask,
-            token_type_ids=token_type_ids,
-            return_dict=False,
-        )
+        self.l1 = BertModel.from_pretrained(bert_path)
+        self.l2 = torch.nn.Dropout(0.3)
+        self.l3 = torch.nn.Linear(768, 6)
+    
+    def forward(self, ids, mask, token_type_ids):
+        _, output_1= self.l1(ids, attention_mask = mask, token_type_ids = token_type_ids)
+        output_2 = self.l2(output_1)
+        output = self.l3(output_2)
         return output
 
 # User selects model for front-end.