Update app.py

app.py

@@ -25,9 +25,13 @@ def get_model_tokenizer_linear():
     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}
-
+    
+@st.cache()
+def predict_topics(top_k):
+    return get_answer_with_desc(text, top_k=top_k, **get_model_tokenizer_linear())[:top_k + 1]
+    
 if len(text) == 1:
       st.markdown("Input is empty, write something!")  
 else:
-    for ms in get_answer_with_desc(text, top_k=top_k, **get_model_tokenizer_linear()):
+    for ms in predict_topics(top_k):
         st.markdown("#### " + ms)