xdgghge

app.py

@@ -2,10 +2,16 @@ import streamlit as st
 from streamlit.elements.altair import generate_chart
 from transformers import AutoTokenizer, AutoModelForCausalLM
 from transformers import pipeline
+
+@st.cache
+def load_model():
+  model_ckpt = "flax-community/gpt2-rap-lyric-generator"
+  tokenizer = AutoTokenizer.from_pretrained(model_ckpt,from_flax=True)
+  model = AutoModelForCausalLM.from_pretrained(model_ckpt,from_flax=True)
+
+
 title = st.title("Loading model")
-model_ckpt = "flax-community/gpt2-rap-lyric-generator"
-tokenizer = AutoTokenizer.from_pretrained(model_ckpt,from_flax=True)
-model = AutoModelForCausalLM.from_pretrained(model_ckpt,from_flax=True)
+model = load_model()
 text_generation = pipeline("text-generation", model=model, tokenizer=tokenizer)
 title.title("Rap lyrics generator")
 artist = st.text_input("Enter the artist", "Wu-Tang Clan")