Update app.py

app.py

@@ -66,22 +66,6 @@ infer_dataset = ToxicityDataset(tweet_df, tokenizer, MAX_LENGTH)
 infer_params = {"batch_size": INFER_BATCH_SIZE, "shuffle": False}
 infer_loader = DataLoader(infer_dataset, **infer_params)
 
-# Have data for BertClass ready for both models
-class BertClass(torch.nn.Module):
-  def __init__(self):
-    super(BertClass, self).__init__()
-    self.l1 = torch.load("pytorch_bert_toxic.bin", map_location=torch.device(device))
-    self.dropout = torch.nn.Dropout(HEAD_DROP_OUT)
-    self.classifier = torch.nn.Linear(768, 6)
-    
-  def forward(self, input_ids, attention_mask, token_type_ids):
-    output_1 = self.l1(input_ids=input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids)
-    hidden_state = output_1[0]
-    pooler = hidden_state[:, 0]
-    pooler = self.dropout(pooler)
-    output = self.classifier(pooler)
-    return output
-
 class PretrainedBertClass(torch.nn.Module):
   def __init__(self):
     super(PretrainedBertClass, self).__init__()
@@ -96,7 +80,7 @@ option = st.selectbox("Select a text analysis model:", ("BERT", "Fine-tuned BERT
 if option == "BERT":
     model = PretrainedBertClass()
 else:
-    model = BertClass()
+    model = torch.load("pytorch_bert_toxic.bin")
 
 # Freeze model and input tokens
 def inference():