Update app.py

app.py

@@ -8,7 +8,7 @@ from sklearn import metrics
 from torch.utils.data import Dataset as set, DataLoader as DL
 from torch import cuda
 import streamlit as st
-from transformers import DistilBertTokenizer as DBT, DistilBertModel as DBM
+from transformers import BertTokenizer as BT, BertModel as BM
 
 # Defined variables for later use
 MAX_LEN = 128
@@ -16,8 +16,7 @@ TRAIN_BATCH_SIZE = 4
 VALID_BATCH_SIZE = 4
 LEARNING_RATE = 5e-05
 
-modName = 'distilbert-base-uncased' # Pre-trained model
-
+modName = 'bert-base-uncased' # Pre-trained model
 
 categories = ['toxic', 'severe_toxic', 'obscene', 'threat', 'insult', 'identity_hate'] # Labels
 
@@ -44,7 +43,7 @@ new = pd.DataFrame()
 new['text'] = data['comment_text']
 new['labels'] = data.iloc[:,1].values.tolist()
 
-tokenizer = DBT.from_pretrained(modName, truncation=True, do_lower_case=True)
+tokenizer = BT.from_pretrained(modName, truncation=True, do_lower_case=True)
 
 class MultiLabelDataset(set):
     def __init__(self, df, tokenizer, max_len):
@@ -102,7 +101,7 @@ for dat in testing_loader:
 class DistilBERTClass(TNN.Module):
     def __init__(self):
         super(DistilBERTClass, self).__init__()
-        self.l1 = DBM.from_pretrained(modName)
+        self.l1 = BM.from_pretrained(modName)
         self.pre_classifier = TNN.Linear(768, 768)
         self.dropout = TNN.Dropout(0.1)
         self.classifier = TNN.Linear(768, 6)