bug fix

app.py

@@ -4,7 +4,7 @@ import streamlit as st
 import os
 from dotenv import load_dotenv
 from views.intro_screen import welcome_screen
-from views.questions_screen import questions_screen
+from views.questions_screen import questions_screen, display_completion_message
 from views.continue_survey import continue_survey_screen
 from css.layout import custom_css
 st.set_page_config(layout="wide")
@@ -81,7 +81,7 @@ def ui():
                 if st.session_state.current_index >= len(data):
                     # If all questions have been answered, show the exit screen
                     print("survey completed")
-                    survey_completed()
+                    display_completion_message()
                 # Otherwise, show questions from where they left off
                 questions_screen(data)
         else:

requirements.txt

@@ -2,4 +2,5 @@ streamlit
 pandas
 python-dotenv
 firebase-admin
-pydantic
+pydantic
+datasets

views/nav_buttons.py

@@ -59,12 +59,26 @@ def navigation_buttons(data, response: Response):
                 st.warning("Please provide all ratings before proceeding.")
             else:
                 if current_index < len(data) - 1:
-                    st.session_state.previous_ratings[data.iloc[st.session_state.current_index]['config_id']] = response.model_ratings
+                    config_id = data.iloc[st.session_state.current_index]['config_id']
+                    if isinstance(response.model_ratings, dict):
+                        st.session_state.previous_ratings[config_id] = {
+                            model: {
+                                'query_v_ratings': ratings.query_v_ratings,
+                                'query_p0_ratings': ratings.query_p0_ratings,
+                                'query_p1_ratings': ratings.query_p1_ratings
+                            }
+                            for model, ratings in response.model_ratings.items()
+                        }
+                    else:
+                        st.session_state.previous_ratings[config_id] = {
+                            model: {
+                                'query_v_ratings': getattr(ratings, 'query_v_ratings', {}),
+                                'query_p0_ratings': getattr(ratings, 'query_p0_ratings', {}),
+                                'query_p1_ratings': getattr(ratings, 'query_p1_ratings', {})
+                            }
+                            for model, ratings in response.model_ratings.items()
+                        }
                     st.session_state.current_index += 1
-                    st.rerun()
-                else:
-                    if st.button("Finish"):
-                        submit_feedback(current_index)
 
     with col3:  # Submit button
         if st.button("Submit & Resume Later"):

views/questions_screen.py

@@ -85,7 +85,7 @@ def render_single_rating(
             options=options,
             format_func=format_func,
             key=f"{key_prefix}",
-            index=stored_rating if stored_rating is not None else 0,
+            index=stored_rating if stored_rating is not None else None,
         )