5

app.py

@@ -9,7 +9,7 @@ SAVE_PATH = "finetuned_deepseek_math"
 
 @st.cache_resource(show_spinner=False)
 def load_model():
-    # # 4-bit quantization configurationi
+    # 4-bit quantization configuration (currently commented out)
     # bnb_config = BitsAndBytesConfig(
     #     load_in_4bit=True,
     #     bnb_4bit_quant_type="nf4",
@@ -63,11 +63,17 @@ def generate_output(prompt, tokenizer, model):
 st.title("Deepseek Math Fine-Tuned Model Inference")
 st.write("Enter your prompt below:")
 
+# Cache the model in Streamlit's session state so it's loaded only once.
+if "tokenizer" not in st.session_state or "model" not in st.session_state:
+    st.session_state.tokenizer, st.session_state.model = load_model()
+
 user_input = st.text_input("Prompt", "🚗 + 🚗 + 🚗 + 🚗 = 20 → 🚗 =")
 
 if st.button("Generate Output"):
     with st.spinner("Generating answer..."):
-        tokenizer, model = load_model()
+        # Use the cached model from session state
+        tokenizer = st.session_state.tokenizer
+        model = st.session_state.model
         output = generate_output(user_input, tokenizer, model)
     st.success("Output generated!")
     st.write("**Input:**", user_input)