Update app.py

app.py

@@ -3,10 +3,26 @@ import pandas as pd
 import torch
 from torch import nn
 from torch.utils.data import Dataset, DataLoader
-from transformers import AutoTokenizer, BertForSequenceClassification
+from transformers import AutoTokenizer, BertModel, BertForSequenceClassification
 from sklearn import metrics
 import streamlit as st
 
+# Have data for BertClass ready for our tuned model.
+class BertClass(torch.nn.Module):
+  def __init__(self):
+    super(BertClass, self).__init__()
+    self.l1 = BertModel.from_pretrained(model_path)
+    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
+
 # Define models to be used
 bert_path = "bert-base-uncased"
 bert_tokenizer = AutoTokenizer.from_pretrained(bert_path)