Update app.py

app.py

@@ -63,7 +63,7 @@ class ToxicityDataset(Dataset):
 
 # Based on user model selection, prepare Dataset and DataLoader
 infer_dataset = ToxicityDataset(tweet_df, tokenizer, MAX_LENGTH)
-infer_params = {"batch_size": INFER_BATCH_SIZE, "shuffle": False}
+infer_params = {"batch_size": INFER_BATCH_SIZE, "shuffle": True}
 infer_loader = DataLoader(infer_dataset, **infer_params)
 
 class BertClass(torch.nn.Module):
@@ -123,9 +123,10 @@ def inference():
             final_outputs.extend(torch.sigmoid(outputs).cpu().detach().numpy().tolist())
     return final_outputs, final_targets
 
+# Get predictions!
 prediction, targets = inference()
-st.write(prediction)
 
+# Format and present findings.
 best_preds = []
 best_labels = []
 for example in prediction:
@@ -135,7 +136,7 @@ for example in prediction:
     best_preds.append(best_prediction)
     best_labels.append(best_label)
 
-st.write("result")
+st.write("Toxicity Classification Result:")
 display_table = tweets_raw[["comment_text"]]
 display_table["Toxicity Classification"] = best_labels
 display_table["Probability"] = best_preds