x

app.py

@@ -5,11 +5,11 @@ from streamlit_chat import message
 
 checkpoint = "."
 
-@st.cache_resource
+tokenizer = AutoTokenizer.from_pretrained(checkpoint)    
+@st.cache
 def get_model():
     model = AutoModelForCausalLM.from_pretrained(checkpoint)
-    tokenizer = AutoTokenizer.from_pretrained(checkpoint)
-    return model, tokenizer
+    return model
      
 st.title("Chat with myGPT 🦄")     
 st.write("This is a LLM that was fine-tuned on a dataset of daily conversations.")
@@ -37,7 +37,7 @@ def submit():
 # prompt = "How long will it take for the poc to finish?"
 # inputs = tokenizer(prompt, return_tensors="pt")
 
-model, tokenizer = get_model()
+model = get_model()
 generation_config = GenerationConfig(max_new_tokens=32, 
                         num_beams=4, 
                         early_stopping=True,