Update utils.py

utils.py

@@ -1,3 +1,5 @@
+import streamlit as st
+
 import torch
 import numpy as np
 
@@ -8,7 +10,7 @@ def get_most_probability_terms(logits, top_k=5):
   _, indices = torch.sort(logits, dim=1, descending=True)
   return indices[:, :top_k]
 
-
+@st.cache(suppress_st_warning=True)
 def predict(text, model, tokenizer, my_linear, top_k=3):
   tokens = tokenizer.encode(text)
   with torch.no_grad():
@@ -18,9 +20,9 @@ def predict(text, model, tokenizer, my_linear, top_k=3):
 
 
 def get_answer_with_desc(text, model, tokenizer, my_linear, top_k=3):
-  codes = predict(text, model, tokenizer, my_linear, top_k=top_k)
+  codes = predict(text, model, tokenizer, my_linear, top_k=10)
   answer = ["Possible text topics:"]
-  for code in codes:
+  for code in codes[:top_k]:
     answer += [code + ": " + classes_desc[code]]
 
   return answer