Update app.py

app.py

@@ -7,6 +7,31 @@ from transformers import AutoTokenizer, BertModel
 from sklearn import metrics
 import streamlit as st
 
+# Define models to be used
+bert_path = "bert-base-uncased"
+bert_tokenizer = AutoTokenizer.from_pretrained(bert_path)
+bert_model = BertForSequenceClassification.from_pretrained(bert_path, num_labels=6)
+tuned_model = model = torch.load("pytorch_bert_toxic.bin")
+
+# Read and format data.
+tweets_raw = pd.read_csv("test.csv", nrows=20)
+labels_raw = pd.read_csv("test_labels.csv", nrows=20)
+
+label_set = ["toxic", "severe_toxic", "obscene", "threat", "insult", "identity_hate"]
+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"))
+if option == "BERT":
+    tokenizer = bert_tokenizer
+    model = bert_model
+else:
+    tokenizer = bert_tokenizer
+    model = tuned_model
+
 # Dataset for loading tables into DataLoader
 class ToxicityDataset(Dataset):
     def __init__(self, dataframe, tokenizer, max_len):
@@ -42,13 +67,20 @@ class ToxicityDataset(Dataset):
             "targets": torch.tensor(self.targets[index], dtype=torch.float),
         }
 
+# Based on user model selection, prepare Dataset and DataLoader
+MAX_LENGTH = 100
+TEST_BATCH_SIZE = 128
+infer_dataset = ToxicityDataset(tweet_df, tokenizer, MAX_LENGTH)
+infer_params = {"batch_size": TEST_BATCH_SIZE, "shuffle": True, "num_workers": 0}
+infer_loader = DataLoader(test_dataset, **test_params)
+
 # Freeze model and input tokens
 def inference():
     model.eval()
     final_targets = []
     final_outputs = []
     with torch.no_grad():
-        for _, data in enumerate(testing_loader, 0):
+        for _, data in enumerate(infer_loader, 0):
             ids = data["ids"].to(device, dtype=torch.long)
             mask = data["mask"].to(device, dtype=torch.long)
             token_type_ids = data["token_type_ids"].to(device, dtype=torch.long)
@@ -58,38 +90,6 @@ def inference():
             final_outputs.extend(torch.sigmoid(outputs).cpu().detach().numpy().tolist())
     return final_outputs, final_targets
 
-# Define models to be used
-bert_path = "bert-base-uncased"
-bert_tokenizer = AutoTokenizer.from_pretrained(bert_path)
-bert_model = BertForSequenceClassification.from_pretrained(bert_path, num_labels=6)
-tuned_model = model = torch.load("pytorch_bert_toxic.bin")
-
-# Read and format data.
-tweets_raw = pd.read_csv("test.csv", nrows=20)
-labels_raw = pd.read_csv("test_labels.csv", nrows=20)
-
-label_set = ["toxic", "severe_toxic", "obscene", "threat", "insult", "identity_hate"]
-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"))
-if option == "BERT":
-    tokenizer = bert_tokenizer
-    model = bert_model
-else:
-    tokenizer = bert_tokenizer
-    model = tuned_model
-
-# Based on user selection, prepare DataLoader
-MAX_LENGTH = 100
-TEST_BATCH_SIZE = 128
-infer_dataset = ToxicityDataset(tweet_df, tokenizer, MAX_LENGTH)
-infer_params = {"batch_size": TEST_BATCH_SIZE, "shuffle": True, "num_workers": 0}
-testing_loader = DataLoader(test_dataset, **test_params)
-
 prediction, targets = inference()
 prediction = np.array(prediction) >= 0.5
 targets = np.argmax(targets, axis=1)