Update app.py

app.py

@@ -7,6 +7,7 @@ from transformers import AutoTokenizer, BertModel
 from sklearn import metrics
 import streamlit as st
 
+# Dataset for loading tables into DataLoader
 class ToxicityDataset(Dataset):
     def __init__(self, dataframe, tokenizer, max_len):
         self.tokenizer = tokenizer
@@ -41,7 +42,7 @@ class ToxicityDataset(Dataset):
             "targets": torch.tensor(self.targets[index], dtype=torch.float),
         }
 
-
+# Freeze model and input tokens
 def inference():
     model.eval()
     final_targets = []
@@ -57,22 +58,23 @@ 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_set = ["toxic", "severe_toxic", "obscene", "threat", "insult", "identity_hate"]
+label_vector = labels_raw[label_set].values.tolist()
 
-MAX_LENGTH = 100
-TEST_BATCH_SIZE = 128
-test_dataset = ToxicityDataset(test_df, tokenizer, MAX_LENGTH)
-test_params = {"batch_size": TEST_BATCH_SIZE, "shuffle": True, "num_workers": 0}
-testing_loader = DataLoader(test_dataset, **test_params)
+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
@@ -81,20 +83,22 @@ 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)
 prediction = np.argmax(prediction, axis=1)
 accuracy = metrics.accuracy_score(targets, prediction)
 f1_score_micro = metrics.f1_score(targets, prediction, average="micro")
 f1_score_macro = metrics.f1_score(targets, prediction, average="macro")
 
-print(f"Accuracy Score = {accuracy}")
-print(f"F1 Score (Micro) = {f1_score_micro}")
-print(f"F1 Score (Macro) = {f1_score_macro}")
-
-# Write results
-st.write("Classification Probabilities")
-st.write(f"{neutralProb:.4f} - NEUTRAL")
-st.write(f"{toxicProb:.4f} - TOXIC")
+st.write(prediction)
+st.write(f"Accuracy Score = {accuracy}")
+st.write(f"F1 Score (Micro) = {f1_score_micro}")
+st.write(f"F1 Score (Macro) = {f1_score_macro}")