Update app.py

app.py

@@ -60,20 +60,24 @@ def search_data(query, search_criteria, search_filter):
     elif search_criteria == 'GSS Tags':
         return d4
 
+def show_data(var):
+    st.session_state['current_var'] = var
+
+def show_search():
+    st.session_state['current_var'] = None
+
 # Main function to run the Streamlit app
 def main():
-    # Page setup
-    st.set_page_config(page_title="AI-Augmented Surveys: Retrodiction Browser", layout="wide")
-    st.title("AI-Augmented Surveys: Retrodiction Browser")
+    st.set_page_config(page_title="AI-Augmented Surveys: Browse Retrodictions", layout="wide")
+    st.title("AI-Augmented Surveys: Browse Retrodictions")
     st.markdown("Junsol Kim, Byungkyu Lee<br />Paper: https://arxiv.org/abs/2305.09620", unsafe_allow_html=True)    
     n_var_per_page = 25
-    N_cards_per_row = 5
+    n_cards_per_row = 5
     current_variable = None
+    if 'current_var' not in st.session_state:
+        st.session_state['current_var'] = None
 
-    if 'search_mode' not in st.session_state:
-        st.session_state['search_mode'] = True
-
-    if st.session_state['search_mode']:
+    if st.session_state['current_var'] is None:
         col1, col2, col3 = st.columns([3, 1, 1])
         with col1:
             search_query = st.text_input('Search available GSS variables')
@@ -101,26 +105,24 @@ def main():
                     st.markdown(f'Page {page} of {maxpage}')
                 button_list = []
                 for n_row, row in results.loc[(page-1)*n_var_per_page:page*n_var_per_page-1].reset_index().iterrows():
-                    i = n_row % N_cards_per_row
+                    i = n_row % n_cards_per_row
                     if i == 0:
                         st.write("---")
-                        cols = st.columns(N_cards_per_row, gap="large")
+                        cols = st.columns(n_cards_per_row, gap="large")
                     # draw the card
-                    with cols[n_row % N_cards_per_row]:
+                    with cols[n_row % n_cards_per_row]:
                         st.markdown(f"**{row['var_name'].strip()}**", unsafe_allow_html=True)
                         st.markdown(f"*{row['var_description'].strip()}*", unsafe_allow_html=True)
                         st.markdown(f"Survey question: {row['question']}", unsafe_allow_html=True)
                         st.markdown(f"GSS tags: {row['subject']}", unsafe_allow_html=True)
-                        button_list.append(st.button('View this variable', key=n_row))
+                        button_list.append(st.button('View this variable', key=row['var_name'], on_click=show_data, args=(row['var_name'], )))
                 for i in range(len(button_list)):
                     if button_list[i]:
                         current_variable = str(results.loc[i, 'var_name'])
                         st.session_state['search_mode'] = False
     else:
-        goback = st.button('Go back')
-        if goback:
-            st.session_state['search_mode'] = True
-        var = current_variable
+        goback = st.button('Go back', on_click=show_search)
+        var = st.session_state['current_var']
         tab_pred = dt_summary1.loc[(dt_summary1.variable == var) & (dt_summary1.pred_type == 'rescale')]
         tab_obs = dt_summary1.loc[(dt_summary1.variable == var) & (dt_summary1.pred_type == 'obsbin')]
         tab_obs = tab_obs.merge(tab_pred[['year']], how='right').rename({'mean': 'obs_mean', 'lci': 'obs_lci', 'uci': 'obs_uci'}, axis=1)