Factor out judge results code so that it persists when the submit button is inactivated.

app.py

@@ -508,6 +508,14 @@ def plot_per_judge_overall_scores(df):
     st.pyplot(plt)
 
 
+def get_selected_models_to_streamlit_column_map(st_columns, selected_models):
+    selected_models_to_streamlit_column_map = {
+        model: st_columns[i % len(st_columns)]
+        for i, model in enumerate(selected_models)
+    }
+    return selected_models_to_streamlit_column_map
+
+
 # Main Streamlit App
 def main():
     st.set_page_config(
@@ -577,6 +585,24 @@ def main():
         if "selected_aggregator" not in st.session_state:
             st.session_state["selected_aggregator"] = None
 
+        # Initialize session state for direct assessment judging.
+        if "direct_assessment_overall_score" not in st.session_state:
+            st.session_state.direct_assessment_overall_score = {}
+        if "direct_assessment_judging_df" not in st.session_state:
+            st.session_state.direct_assessment_judging_df = defaultdict(dict)
+        if "direct_assessment_judging_responses" not in st.session_state:
+            st.session_state.direct_assessment_judging_responses = defaultdict(dict)
+        if "direct_assessment_overall_scores" not in st.session_state:
+            st.session_state.direct_assessment_overall_scores = defaultdict(dict)
+        if "judging_status" not in st.session_state:
+            st.session_state.judging_status = "incomplete"
+        if "direct_assessment_config" not in st.session_state:
+            st.session_state.direct_assessment_config = {}
+        if "pairwise_comparison_config" not in st.session_state:
+            st.session_state.pairwise_comparison_config = {}
+        if "assessment_type" not in st.session_state:
+            st.session_state.assessment_type = None
+
         with st.form(key="prompt_form"):
             st.markdown("#### LLM Council Member Selection")
 
@@ -608,9 +634,11 @@ def main():
             # Render the chats.
             response_columns = st.columns(3)
 
-            selected_models_to_streamlit_column_map = {
-                model: response_columns[i] for i, model in enumerate(selected_models)
-            }
+            selected_models_to_streamlit_column_map = (
+                get_selected_models_to_streamlit_column_map(
+                    response_columns, selected_models
+                )
+            )
 
             # Fetching and streaming responses from each selected model
             for selected_model in st.session_state.selected_models:
@@ -643,7 +671,7 @@ def main():
                     selected_aggregator, aggregator_prompt
                 )
                 if aggregator_stream:
-                    st.session_state["responses"]["agg__" + selected_aggregator] = (
+                    st.session_state.responses["agg__" + selected_aggregator] = (
                         message_placeholder.write_stream(aggregator_stream)
                     )
 
@@ -654,10 +682,11 @@ def main():
             st.markdown("#### Responses")
 
             response_columns = st.columns(3)
-            selected_models_to_streamlit_column_map = {
-                model: response_columns[i]
-                for i, model in enumerate(st.session_state.selected_models)
-            }
+            selected_models_to_streamlit_column_map = (
+                get_selected_models_to_streamlit_column_map(
+                    response_columns, st.session_state.selected_models
+                )
+            )
             for response_model, response in st.session_state.responses.items():
                 st_column = selected_models_to_streamlit_column_map.get(
                     response_model, response_columns[0]
@@ -671,106 +700,56 @@ def main():
 
         # Judging.
         if st.session_state.responses_collected:
-            st.markdown("#### Judging Configuration")
+            with st.form(key="judging_form"):
+                st.markdown("#### Judging Configuration")
 
-            # Choose the type of assessment
-            assessment_type = st.radio(
-                "Select the type of assessment",
-                options=["Direct Assessment", "Pairwise Comparison"],
-            )
+                # Choose the type of assessment
+                assessment_type = st.radio(
+                    "Select the type of assessment",
+                    options=["Direct Assessment", "Pairwise Comparison"],
+                )
 
-            _, center_column, _ = st.columns([3, 5, 3])
+                _, center_column, _ = st.columns([3, 5, 3])
+
+                # Depending on the assessment type, render different forms
+                if assessment_type == "Direct Assessment":
+                    # Direct assessment prompt.
+                    with center_column.expander("Direct Assessment Prompt"):
+                        direct_assessment_prompt = st.text_area(
+                            "Prompt for the Direct Assessment",
+                            value=get_default_direct_assessment_prompt(
+                                user_prompt=user_prompt
+                            ),
+                            height=500,
+                            key="direct_assessment_prompt",
+                        )
 
-            # Depending on the assessment type, render different forms
-            if assessment_type == "Direct Assessment":
+                    # TODO: Add option to edit criteria list with a basic text field.
+                    criteria_list = DEFAULT_DIRECT_ASSESSMENT_CRITERIA_LIST
 
-                # Initialize session state for direct assessment.
-                if "direct_assessment_overall_score" not in st.session_state:
-                    st.session_state["direct_assessment_overall_score"] = {}
-                if "direct_assessment_judging_df" not in st.session_state:
-                    st.session_state["direct_assessment_judging_df"] = {}
-                    for response_model in selected_models:
-                        st.session_state["direct_assessment_judging_df"][
-                            response_model
-                        ] = {}
-                    # aggregator model
-                    st.session_state["direct_assessment_judging_df"][
-                        "agg__" + selected_aggregator
-                    ] = {}
-                if "direct_assessment_judging_responses" not in st.session_state:
-                    st.session_state["direct_assessment_judging_responses"] = {}
-                    for response_model in selected_models:
-                        st.session_state["direct_assessment_judging_responses"][
-                            response_model
-                        ] = {}
-                    # aggregator model
-                    st.session_state["direct_assessment_judging_responses"][
-                        "agg__" + selected_aggregator
-                    ] = {}
-                if "direct_assessment_overall_scores" not in st.session_state:
-                    st.session_state["direct_assessment_overall_scores"] = {}
-                    for response_model in selected_models:
-                        st.session_state["direct_assessment_overall_scores"][
-                            response_model
-                        ] = {}
-                    st.session_state["direct_assessment_overall_scores"][
-                        "agg__" + selected_aggregator
-                    ] = {}
-                if "judging_status" not in st.session_state:
-                    st.session_state["judging_status"] = "incomplete"
-
-                # Direct assessment prompt.
-                with center_column.expander("Direct Assessment Prompt"):
-                    direct_assessment_prompt = st.text_area(
-                        "Prompt for the Direct Assessment",
-                        value=get_default_direct_assessment_prompt(
-                            user_prompt=user_prompt
-                        ),
-                        height=500,
-                        key="direct_assessment_prompt",
+                    judging_submit_button = st.form_submit_button(
+                        "Submit Judging", use_container_width=True
                     )
 
-                # TODO: Add option to edit criteria list with a basic text field.
-                criteria_list = DEFAULT_DIRECT_ASSESSMENT_CRITERIA_LIST
-
-                # Create DirectAssessment object when form is submitted
-                if center_column.button(
-                    "Submit Direct Assessment", use_container_width=True
-                ):
-
-                    # Render the chats.
-                    response_columns = st.columns(3)
-                    selected_models_to_streamlit_column_map = {
-                        model: response_columns[i]
-                        for i, model in enumerate(selected_models)
-                    }
-                    for response_model, response in st.session_state[
-                        "responses"
-                    ].items():
-                        st_column = selected_models_to_streamlit_column_map.get(
-                            response_model, response_columns[0]
-                        )
-                        with st_column:
-                            with st.chat_message(
-                                get_ui_friendly_name(response_model),
-                                avatar=get_llm_avatar(response_model),
-                            ):
-                                st.write(get_ui_friendly_name(response_model))
-                                st.write(response)
-
-                    # Submit direct asssessment.
-                    responses_for_judging = st.session_state["responses"]
+            if judging_submit_button:
+                st.session_state.assessment_type = assessment_type
+                st.session_state.direct_assessment_config = {
+                    "prompt": direct_assessment_prompt,
+                    "criteria_list": criteria_list,
+                }
 
-                    response_judging_columns = st.columns(3)
+                responses_for_judging = st.session_state.responses
 
-                    responses_for_judging_to_streamlit_column_map = {
-                        model: response_judging_columns[i % 3]
-                        for i, model in enumerate(responses_for_judging.keys())
-                    }
+                # Get judging responses.
+                response_judging_columns = st.columns(3)
+                responses_for_judging_to_streamlit_column_map = (
+                    get_selected_models_to_streamlit_column_map(
+                        response_judging_columns, responses_for_judging.keys()
+                    )
+                )
 
-                    # Get judging responses.
+                if st.session_state.assessment_type == "Direct Assessment":
                     for response_model, response in responses_for_judging.items():
-
                         st_column = responses_for_judging_to_streamlit_column_map[
                             response_model
                         ]
@@ -792,7 +771,7 @@ def main():
 
                             for judging_model in selected_models:
                                 with st.expander(
-                                    get_ui_friendly_name(judging_model), expanded=False
+                                    get_ui_friendly_name(judging_model), expanded=True
                                 ):
                                     with st.chat_message(
                                         judging_model,
@@ -811,7 +790,6 @@ def main():
                                         )
                             # When all of the judging is finished for the given response, get the actual
                             # values, parsed.
-                            # TODO.
                             judging_responses = st.session_state[
                                 "direct_assessment_judging_responses"
                             ][response_model]
@@ -872,104 +850,192 @@ def main():
                             st.write(f"Overall Score: {overall_score:.2f}")
                             st.write(f"Controversy: {controversy:.2f}")
 
-                    st.session_state["judging_status"] = "complete"
+                    st.session_state.judging_status = "complete"
+            # If judging is complete, but the submit button is cleared, still render the results.
+            elif st.session_state.judging_status == "complete":
+                if st.session_state.assessment_type == "Direct Assessment":
+                    responses_for_judging = st.session_state.responses
 
-                # Judging is complete.
-                # The session state now contains the overall scores for each response from each judge.
-                if st.session_state["judging_status"] == "complete":
-                    st.write("#### Results")
+                    # Get judging responses.
+                    response_judging_columns = st.columns(3)
+                    responses_for_judging_to_streamlit_column_map = (
+                        get_selected_models_to_streamlit_column_map(
+                            response_judging_columns, responses_for_judging.keys()
+                        )
+                    )
 
-                    overall_scores_df_raw = pd.DataFrame(
-                        st.session_state["direct_assessment_overall_scores"]
-                    ).reset_index()
+                    for response_model, response in responses_for_judging.items():
+                        st_column = responses_for_judging_to_streamlit_column_map[
+                            response_model
+                        ]
 
-                    overall_scores_df = pd.melt(
-                        overall_scores_df_raw,
-                        id_vars=["index"],
-                        var_name="response_model",
-                        value_name="score",
-                    ).rename(columns={"index": "judging_model"})
+                        with st_column:
+                            st.write(
+                                f"Judging for {get_ui_friendly_name(response_model)}"
+                            )
+                            judging_prompt = get_direct_assessment_prompt(
+                                direct_assessment_prompt=direct_assessment_prompt,
+                                user_prompt=user_prompt,
+                                response=response,
+                                criteria_list=criteria_list,
+                                options=SEVEN_POINT_DIRECT_ASSESSMENT_OPTIONS,
+                            )
 
-                    # Print the overall winner.
-                    overall_winner = overall_scores_df.loc[
-                        overall_scores_df["score"].idxmax()
-                    ]
+                            with st.expander("Final Judging Prompt"):
+                                st.code(judging_prompt)
 
-                    st.write(
-                        f"**Overall Winner:** {get_ui_friendly_name(overall_winner['response_model'])}"
-                    )
-                    # Find how much the standard deviation overlaps with other models.
-                    # Calculate separability.
-                    # TODO.
-                    st.write(f"**Confidence:** {overall_winner['score']:.2f}")
-
-                    left_column, right_column = st.columns([1, 1])
-                    with left_column:
-                        plot_overall_scores(overall_scores_df)
-
-                    with right_column:
-                        # All overall scores.
-                        overall_scores_df = overall_scores_df[
-                            ["response_model", "judging_model", "score"]
-                        ]
-                        overall_scores_df["response_model"] = overall_scores_df[
-                            "response_model"
-                        ].apply(get_ui_friendly_name)
-                        overall_scores_df["judging_model"] = overall_scores_df[
-                            "judging_model"
-                        ].apply(get_ui_friendly_name)
-
-                        with st.expander("Overall scores from all judges"):
-                            st.dataframe(overall_scores_df)
-
-                    # All criteria scores.
-                    with right_column:
-                        all_scores_df = pd.DataFrame()
-                        for response_model, score_df in st.session_state[
-                            "direct_assessment_judging_df"
-                        ].items():
-                            score_df["response_model"] = response_model
-                            all_scores_df = pd.concat([all_scores_df, score_df])
-                        all_scores_df = all_scores_df.reset_index()
-                        all_scores_df = all_scores_df.drop(columns="index")
-
-                        # Reorder the columns
-                        all_scores_df = all_scores_df[
-                            [
-                                "response_model",
-                                "judging_model",
-                                "criteria",
-                                "score",
-                                "explanation",
-                            ]
-                        ]
-                        all_scores_df["response_model"] = all_scores_df[
-                            "response_model"
-                        ].apply(get_ui_friendly_name)
-                        all_scores_df["judging_model"] = all_scores_df[
-                            "judging_model"
-                        ].apply(get_ui_friendly_name)
+                            for judging_model in selected_models:
+                                with st.expander(
+                                    get_ui_friendly_name(judging_model), expanded=True
+                                ):
+                                    with st.chat_message(
+                                        judging_model,
+                                        avatar=PROVIDER_TO_AVATAR_MAP[judging_model],
+                                    ):
+                                        st.write(
+                                            st.session_state.direct_assessment_judging_responses[
+                                                response_model
+                                            ][
+                                                judging_model
+                                            ]
+                                        )
+                            # When all of the judging is finished for the given response, get the actual
+                            # values, parsed.
+                            judging_responses = (
+                                st.session_state.direct_assessment_judging_responses[
+                                    response_model
+                                ]
+                            )
 
-                        with st.expander(
-                            "Criteria-specific scores and explanations from all judges"
-                        ):
-                            st.dataframe(all_scores_df)
+                            parse_judging_response_prompt = (
+                                get_parse_judging_response_for_direct_assessment_prompt(
+                                    judging_responses,
+                                    criteria_list,
+                                    SEVEN_POINT_DIRECT_ASSESSMENT_OPTIONS,
+                                )
+                            )
 
-            elif assessment_type == "Pairwise Comparison":
-                pass
+                            plot_criteria_scores(
+                                st.session_state.direct_assessment_judging_df[
+                                    response_model
+                                ]
+                            )
+
+                            plot_per_judge_overall_scores(
+                                st.session_state.direct_assessment_judging_df[
+                                    response_model
+                                ]
+                            )
+
+                            grouped = (
+                                st.session_state.direct_assessment_judging_df[
+                                    response_model
+                                ]
+                                .groupby(["judging_model"])
+                                .agg({"score": ["mean"]})
+                                .reset_index()
+                            )
+                            grouped.columns = ["judging_model", "overall_score"]
+
+                            overall_score = grouped["overall_score"].mean()
+                            controversy = grouped["overall_score"].std()
+                            st.write(f"Overall Score: {overall_score:.2f}")
+                            st.write(f"Controversy: {controversy:.2f}")
+
+            # Judging is complete, stuff that would be rendered that's not stream-specific.
+            # The session state now contains the overall scores for each response from each judge.
+            if st.session_state.judging_status == "complete":
+                st.write("#### Results")
+
+                overall_scores_df_raw = pd.DataFrame(
+                    st.session_state["direct_assessment_overall_scores"]
+                ).reset_index()
+
+                overall_scores_df = pd.melt(
+                    overall_scores_df_raw,
+                    id_vars=["index"],
+                    var_name="response_model",
+                    value_name="score",
+                ).rename(columns={"index": "judging_model"})
+
+                # Print the overall winner.
+                overall_winner = overall_scores_df.loc[
+                    overall_scores_df["score"].idxmax()
+                ]
+
+                st.write(
+                    f"**Overall Winner:** {get_ui_friendly_name(overall_winner['response_model'])}"
+                )
+                # Find how much the standard deviation overlaps with other models
+                # TODO: Calculate separability.
+                st.write(f"**Confidence:** {overall_winner['score']:.2f}")
+
+                left_column, right_column = st.columns([1, 1])
+                with left_column:
+                    plot_overall_scores(overall_scores_df)
+
+                with right_column:
+                    # All overall scores.
+                    overall_scores_df = overall_scores_df[
+                        ["response_model", "judging_model", "score"]
+                    ]
+                    overall_scores_df["response_model"] = overall_scores_df[
+                        "response_model"
+                    ].apply(get_ui_friendly_name)
+                    overall_scores_df["judging_model"] = overall_scores_df[
+                        "judging_model"
+                    ].apply(get_ui_friendly_name)
+
+                    with st.expander("Overall scores from all judges"):
+                        st.dataframe(overall_scores_df)
+
+                # All criteria scores.
+                with right_column:
+                    all_scores_df = pd.DataFrame()
+                    for response_model, score_df in st.session_state[
+                        "direct_assessment_judging_df"
+                    ].items():
+                        score_df["response_model"] = response_model
+                        all_scores_df = pd.concat([all_scores_df, score_df])
+                    all_scores_df = all_scores_df.reset_index()
+                    all_scores_df = all_scores_df.drop(columns="index")
+
+                    # Reorder the columns
+                    all_scores_df = all_scores_df[
+                        [
+                            "response_model",
+                            "judging_model",
+                            "criteria",
+                            "score",
+                            "explanation",
+                        ]
+                    ]
+                    all_scores_df["response_model"] = all_scores_df[
+                        "response_model"
+                    ].apply(get_ui_friendly_name)
+                    all_scores_df["judging_model"] = all_scores_df[
+                        "judging_model"
+                    ].apply(get_ui_friendly_name)
+
+                    with st.expander(
+                        "Criteria-specific scores and explanations from all judges"
+                    ):
+                        st.dataframe(all_scores_df)
 
         # Token usage.
-        with st.expander("Token Usage"):
-            st.write("Input tokens used.")
-            st.write(st.session_state.input_token_usage)
-            st.write(
-                f"Input Tokens Total: {sum(st.session_state.input_token_usage.values())}"
-            )
-            st.write("Output tokens used.")
-            st.write(st.session_state.output_token_usage)
-            st.write(
-                f"Output Tokens Total: {sum(st.session_state.output_token_usage.values())}"
-            )
+        if st.session_state.responses:
+            st.divider()
+            with st.expander("Token Usage"):
+                st.write("Input tokens used.")
+                st.write(st.session_state.input_token_usage)
+                st.write(
+                    f"Input Tokens Total: {sum(st.session_state.input_token_usage.values())}"
+                )
+                st.write("Output tokens used.")
+                st.write(st.session_state.output_token_usage)
+                st.write(
+                    f"Output Tokens Total: {sum(st.session_state.output_token_usage.values())}"
+                )
 
     else:
         with cols[1]: