Update app.py

app.py

@@ -5,13 +5,9 @@ import torch.nn as nn
 from transformers import RobertaTokenizer, RobertaModel
 
 @st.cache(suppress_st_warning=True)
-def init_tokenizer():
+def init():
     tokenizer = RobertaTokenizer.from_pretrained("roberta-large-mnli")
-    return tokenizer
 
-
-@st.cache(suppress_st_warning=True)
-def init_model():
     model = RobertaModel.from_pretrained("roberta-large-mnli")
 
     model.pooler = nn.Sequential(
@@ -24,6 +20,8 @@ def init_model():
     
     model_path = 'model.pt'
     model.load_state_dict(torch.load(model_path, map_location=torch.device("cpu")))
+    model.eval()
+    return tokenizer, model
 
 cats = ['Computer Science', 'Economics', 'Electrical Engineering', 
         'Mathematics', 'Physics', 'Biology', 'Finance', 'Statistics']
@@ -38,6 +36,7 @@ def predict(outputs):
             top += percent
             st.write(f'{cat}: {round(percent, 1)}')
 
+tokenizer, model = init()
         
 st.markdown("### Title")
 
@@ -50,10 +49,8 @@ abstract = st.text_area("Enter abstract", height=200)
 if not title:
     st.warning("Please fill out so required fields")
 else:    
-    tokenizer = init_tokenizer()
-    model = init_model()
-
     encoded_input = tokenizer(title + '. ' + abstract, return_tensors='pt', padding=True, 
                           max_length = 512, truncation=True)
-    outputs = model(**encoded_input).pooler_output[:, 0, :]
-    predict(outputs)
+    with torch.no_grad():
+        outputs = model(**encoded_input).pooler_output[:, 0, :]
+        predict(outputs)