Update app.py

app.py

@@ -7,8 +7,14 @@ from transformers import AutoTokenizer, BertModel, BertForSequenceClassification
 from sklearn import metrics
 import streamlit as st
 
-# Define Torch device. Enable CUDA if available.
+# Define constants. Enable CUDA if available.
+MAX_LENGTH = 100
+INFER_BATCH_SIZE = 128
+HEAD_DROP_OUT = 0.4
+
 device = "cuda" if torch.cuda.is_available() else "cpu"
+bert_path = "bert-base-uncased"
+tokenizer = AutoTokenizer.from_pretrained(bert_path)
 
 # Read and format data.
 tweets_raw = pd.read_csv("test.csv", nrows=20)
@@ -56,9 +62,6 @@ 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)
@@ -90,13 +93,9 @@ class PretrainedBertClass(torch.nn.Module):
 
 # 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