Update app.py

app.py

@@ -10,31 +10,6 @@ import streamlit as st
 # Define Torch device. Enable CUDA if available.
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# Have data for BertClass ready for both models
-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
-
-class PretrainedBertClass(torch.nn.Module):
-  def __init__(self):
-    super(BertClass, self).__init__()
-    self.l1 = BertForSequenceClassification.from_pretrained(bert_path, num_labels=6)
-    
-  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 output
-
 # Read and format data.
 tweets_raw = pd.read_csv("test.csv", nrows=20)
 labels_raw = pd.read_csv("test_labels.csv", nrows=20)
@@ -45,17 +20,6 @@ label_vector = labels_raw[label_set].values.tolist()
 tweet_df = tweets_raw[["comment_text"]]
 tweet_df["labels"] = label_vector
 
-# User selects model for front-end.
-option = st.selectbox("Select a text analysis model:", ("BERT", "Fine-tuned BERT"))
-
-bert_path = "bert-base-uncased"
-if option == "BERT":
-    tokenizer = AutoTokenizer.from_pretrained(bert_path)
-    model = PretrainedBertClass()
-else:
-    tokenizer = AutoTokenizer.from_pretrained(bert_path)
-    model = torch.load("pytorch_bert_toxic.bin", map_location=torch.device(device))
-
 # Dataset for loading tables into DataLoader
 class ToxicityDataset(Dataset):
     def __init__(self, dataframe, tokenizer, max_len):
@@ -94,10 +58,47 @@ class ToxicityDataset(Dataset):
 # Based on user model selection, prepare Dataset and DataLoader
 MAX_LENGTH = 100
 INFER_BATCH_SIZE = 128
+HEAD_DROP_OUT = 0.4
 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__()
+    self.l1 = BertForSequenceClassification.from_pretrained(bert_path, num_labels=6)
+    
+  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 output
+
+# User selects model for front-end.
+option = st.selectbox("Select a text analysis model:", ("BERT", "Fine-tuned BERT"))
+
+bert_path = "bert-base-uncased"
+if option == "BERT":
+    tokenizer = AutoTokenizer.from_pretrained(bert_path)
+    model = PretrainedBertClass()
+else:
+    tokenizer = AutoTokenizer.from_pretrained(bert_path)
+    model = BertClass()
+
 # Freeze model and input tokens
 def inference():
     model.eval()