Update app.py

app.py

@@ -87,11 +87,22 @@ st.sidebar.write("📥 Load GSM8K")
 gsm8k_data = load_gsm8k_dataset()
 st.sidebar.write("📊 GSM8K loaded:", len(gsm8k_data), "samples")
 
-question_index = st.selectbox("🔢 Select GSM8K question index", range(len(gsm8k_data)))
+# Check for random question index in query params
+random_index = st.query_params.get("question_index")
+if random_index is not None:
+    try:
+        default_index = int(random_index)
+    except (ValueError, TypeError):
+        default_index = 0
+else:
+    default_index = 0
+
+question_index = st.selectbox("🔢 Select GSM8K question index", range(len(gsm8k_data)), index=default_index)
 
 if st.button("🎲 Pick Random Question"):
-    question_index = random.randint(0, len(gsm8k_data) - 1)
-    st.query_params.update(question_index=question_index)
+    new_random_index = random.randint(0, len(gsm8k_data) - 1)
+    st.query_params.update(question_index=new_random_index)
+    st.rerun()  # Force app to rerun to update the selectbox
 
 default_prompt = "Jasper has 5 apples and eats 2 of them. How many apples does he have left?"
 selected_question = gsm8k_data[question_index]["question"] if question_index is not None else default_prompt
@@ -111,6 +122,9 @@ prompt = st.text_area("Enter your math prompt:", selected_question)
 
 # Generation
 if st.button("Generate Response", key="manual"):
+    # Check if the current prompt is from GSM8K dataset
+    is_gsm8k_question = prompt == selected_question
+    
     with st.sidebar:
         with st.spinner("🔄 Generating..."):
 
@@ -143,5 +157,8 @@ if st.button("Generate Response", key="manual"):
         st.write(input_text)
     st.subheader("🧠 Model Output")
     st.success(response_only)
-    st.subheader("✅ Correct Answer (GSM8K)")
-    st.info(correct_answer)
+    
+    # Only show correct answer if using actual GSM8K question
+    if is_gsm8k_question:
+        st.subheader("✅ Correct Answer (GSM8K)")
+        st.info(correct_answer)