reorganize model loading

app.py

@@ -4,14 +4,13 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 # model_name = "flax-community/gpt-code-clippy-1.3B-apps-alldata"
 model_name = "flax-community/gpt-code-clippy-125M-apps-alldata"
 
-@st.cache(allow_output_mutation=True)
+@st.cache(allow_output_mutation=True, max_entries=1)
 def get_model():
-    return AutoModelForCausalLM.from_pretrained(model_name)
-
-def get_tokenizer():
+    model = AutoModelForCausalLM.from_pretrained(model_name)
     tokenizer = AutoTokenizer.from_pretrained(model_name)
     tokenizer.pad_token = tokenizer.eos_token
-    return tokenizer
+    return (model, tokenizer)
+
 
 def format_input(question, starter_code=""):
     answer_type = "\nUse Call-Based format\n" if starter_code else \
@@ -30,6 +29,7 @@ def generate_solution(model, tokenizer, question, starter_code="", temperature=1
         do_sample=True,
         top_p=0.95,
         pad_token_id=tokenizer.pad_token_id,
+        eos_token_id=tokenizer.eos_token_id,
         early_stopping=True,
         temperature=temperature,
         num_beams=int(num_beams),
@@ -105,8 +105,7 @@ def run():
     )
 
     # main body
-    model = get_model()
-    tokenizer = get_tokenizer()
+    model, tokenizer = get_model()
 
     question = st.text_input(
         "Problem: ",
@@ -121,7 +120,7 @@ def run():
     submit_button = st.button("Solve")
 
     if submit_button:
-        
+        st.text("Solution:")
         output = generate_solution(model, tokenizer, question, starter_code, temperature, num_beams)
         st.code(output, language="python")