Update app.py

app.py

@@ -48,11 +48,26 @@ def load_model():
     tokenizer = LlamaTokenizer.from_pretrained(repo_id)
     return model, tokenizer
 
-if st.button("Load Model"):
-    model1, tokenizer1 = load_model()
-    model_loaded = True
-
 print ("Model Loaded")
+# Initialize session state variables
+if "model_loaded" not in st.session_state:
+    st.session_state["model_loaded"] = False
+if "model" not in st.session_state:
+    st.session_state["model"] = None
+if "tokenizer" not in st.session_state:
+    st.session_state["tokenizer"] = None
+
+# Display the "Load Model" button
+if not st.session_state["model_loaded"]:
+    if st.button("Load Model"):
+        model1, tokenizer1 = load_model(repo_id)
+        st.session_state["model"] = model1
+        st.session_state["tokenizer"] = tokenizer1
+        st.session_state["model_loaded"] = True
+else:
+    model1 = st.session_state["model"]
+    tokenizer1 = st.session_state["tokenizer"]
+
 
 if model_loaded:
 # Set up initial values for pipeline parameters