ADD: Caching of model on streamlit

app.py

@@ -15,10 +15,10 @@ def main():
     st.title('Emoji-motion!')
 
     example_prompts = [
+        "This space is lit!!",
         "Today is going to be awesome!",
-        "Pity those who don't feel anything at all.",
-        "I envy people that know love.",
-        "Nature is so beautiful"]
+        "I love Machine Learning",
+        "Cool cool cool no doubt no doubt no doubt"]
 
     example = st.selectbox("Choose a pre-defined example", example_prompts)
 
@@ -31,9 +31,6 @@ def main():
     BASE_MODEL = st.selectbox("Choose a model", models_to_choose)
     TOP_N = 5
 
-    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
-    model = AutoModelForSequenceClassification.from_pretrained(BASE_MODEL)
-
     def preprocess(text):
         new_text = []
         for t in text.split(" "):
@@ -42,17 +39,20 @@ def main():
             new_text.append(t)
         return " ".join(new_text)
 
+    @st.cache(allow_output_mutation=True, suppress_st_warning=True, show_spinner=False)
+    def load_model():
+        tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
+        model = AutoModelForSequenceClassification.from_pretrained(BASE_MODEL)
+        return model, tokenizer
+
     def get_top_emojis(text, top_n=TOP_N):
+        model, tokenizer = load_model()
         preprocessed = preprocess(text)
         inputs = tokenizer(preprocessed, return_tensors="pt")
         preds = model(**inputs).logits
         scores = torch.nn.functional.softmax(preds, dim=-1).detach().numpy()
         ranking = np.argsort(scores)
-        print(ranking)
         ranking = ranking.squeeze()[::-1][:top_n]
-        print(scores)
-        print(ranking)
-        print(model.config.id2label)
         emojis = [model.config.id2label[i] for i in ranking]
         return ', '.join(map(str, emojis))