Update app.py

app.py

@@ -16,12 +16,15 @@ abstract = st.text_area("Enter the abstract of the article", height=350)
 top_k = st.slider('How many topics from top to show?', 1, 10, 3)
 text = title + " " + abstract
 
-tokenizer = DistilBertTokenizer.from_pretrained("distilbert-base-cased")
-model = DistilBertModel.from_pretrained("distilbert-base-cased")
-
-n_classes = 40
-my_linear = nn.Linear(in_features=768, out_features=n_classes, bias=True)
-my_linear.load_state_dict(torch.load(MY_LINEAR_NAME, map_location=torch.device('cpu')))
-
-for ms in get_answer_with_desc(text, model, tokenizer, my_linear, top_k=top_k):
+@st.cache(suppress_st_warning=True)
+def get_model_tokenizer_linear():
+    tokenizer = DistilBertTokenizer.from_pretrained("distilbert-base-cased")
+    model = DistilBertModel.from_pretrained("distilbert-base-cased")
+    
+    n_classes = 40
+    my_linear = nn.Linear(in_features=768, out_features=n_classes, bias=True)
+    my_linear.load_state_dict(torch.load(MY_LINEAR_NAME, map_location=torch.device('cpu')))
+    return {"model": model, "tokenizer": tokenizer, "my_linear": my_linear}
+
+for ms in get_answer_with_desc(text, top_k=top_k, **get_model_tokenizer_linear()):
     st.markdown("#### " + ms)