xdgghge

app.py

@@ -8,10 +8,11 @@ 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)
+  return tokenizer, model
 
 
 title = st.title("Loading model")
-model = load_model()
+tokenizer, 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")