Overall scores graph complete.

.gitignore

@@ -1,2 +1,3 @@
 env/
-client_secret.json
+client_secret.json
+__pycache__

app.py

@@ -15,10 +15,15 @@ from constants import (
     LLM_TO_UI_NAME_MAP,
 )
 from prompts import *
-from judging_dataclasses import DirectAssessmentJudgingResponse
+from judging_dataclasses import (
+    DirectAssessmentJudgingResponse,
+    DirectAssessmentCriterionScore,
+    DirectAssessmentCriteriaScores,
+)
 import pandas as pd
 import seaborn as sns
 import matplotlib.pyplot as plt
+import numpy as np
 
 dotenv.load_dotenv()
 
@@ -67,6 +72,16 @@ def anthropic_streamlit_streamer(stream):
                 break  # End of message, stop streaming
 
 
+def get_ui_friendly_name(llm):
+    if "agg__" in llm:
+        return (
+            "MoA ("
+            + LLM_TO_UI_NAME_MAP.get(llm.split("__")[1], llm.split("__")[1])
+            + ")"
+        )
+    return LLM_TO_UI_NAME_MAP.get(llm, llm)
+
+
 def google_streamlit_streamer(stream):
     for chunk in stream:
         yield chunk.text
@@ -146,22 +161,6 @@ def get_llm_response_stream(model_identifier, prompt):
         return None
 
 
-def get_response_key(model):
-    return model + "__response"
-
-
-def get_model_from_response_key(response_key):
-    return response_key.split("__")[0]
-
-
-def get_direct_assessment_judging_key(judge_model, response_model):
-    return "direct_assessment_judge__" + judge_model + "__" + response_model
-
-
-def get_aggregator_response_key(model):
-    return model + "__aggregator_response"
-
-
 def create_dataframe_for_direct_assessment_judging_response(
     response: DirectAssessmentJudgingResponse,
 ):
@@ -203,21 +202,6 @@ def render_criteria_form(criteria_num):
     )
 
 
-def get_response_mapping():
-    # Inspect the session state for all the responses.
-    # This is a dictionary mapping model names to their responses.
-    # The aggregator response is also included in this mapping under the key "<model>__aggregator_response".
-    response_mapping = {}
-    for key in st.session_state.keys():
-        if "judge" in key:
-            continue
-        if key.endswith("__response"):
-            response_mapping[get_model_from_response_key(key)] = st.session_state[key]
-        if key.endswith("__aggregator_response"):
-            response_mapping[key] = st.session_state[key]
-    return response_mapping
-
-
 def format_likert_comparison_options(options):
     return "\n".join([f"{i + 1}: {option}" for i, option in enumerate(options)])
 
@@ -252,7 +236,7 @@ def get_default_direct_assessment_prompt(user_prompt):
 def get_aggregator_prompt(aggregator_prompt, user_prompt, llms):
     responses_from_other_llms = "\n\n".join(
         [
-            f"{get_ui_friendly_name(model)} START\n{st.session_state.get(get_response_key(model))}\n\n{get_ui_friendly_name(model)} END\n\n\n"
+            f"{get_ui_friendly_name(model)} START\n{st.session_state['responses'][model]}\n\n{get_ui_friendly_name(model)} END\n\n\n"
             for model in llms
         ]
     )
@@ -270,10 +254,6 @@ def get_default_aggregator_prompt(user_prompt, llms):
     )
 
 
-def get_ui_friendly_name(llm):
-    return LLM_TO_UI_NAME_MAP.get(llm, llm)
-
-
 def get_parse_judging_response_for_direct_assessment_prompt(
     judging_responses: dict[str, str],
     criteria_list,
@@ -292,34 +272,58 @@ def get_parse_judging_response_for_direct_assessment_prompt(
     )
 
 
-def get_model_from_direct_assessment_judging_key(judging_key):
-    return judging_key.split("__")[1]
-
-
-def get_direct_assessment_judging_responses():
-    # Get the judging responses from the session state.
-    judging_responses = {}
-    for key in st.session_state.keys():
-        if key.startswith("direct_assessment_judge__"):
-            judging_responses[get_model_from_direct_assessment_judging_key(key)] = (
-                st.session_state[key]
-            )
-    return judging_responses
-
-
-def parse_judging_responses(prompt: str) -> DirectAssessmentJudgingResponse:
-    completion = client.beta.chat.completions.parse(
-        model="gpt-4o-mini",
-        messages=[
-            {
-                "role": "system",
-                "content": "Parse the judging responses into structured data.",
-            },
-            {"role": "user", "content": prompt},
-        ],
-        response_format=DirectAssessmentJudgingResponse,
-    )
-    return completion.choices[0].message.parsed
+DEBUG_MODE = True
+
+
+def parse_judging_responses(
+    prompt: str, judging_responses: dict[str, str]
+) -> DirectAssessmentJudgingResponse:
+    if DEBUG_MODE:
+        return DirectAssessmentJudgingResponse(
+            judging_models=[
+                DirectAssessmentCriteriaScores(
+                    model="together://meta-llama/Meta-Llama-3.1-8B-Instruct-Turbo",
+                    criteria_scores=[
+                        DirectAssessmentCriterionScore(
+                            criterion="helpfulness", score=3, explanation="explanation1"
+                        ),
+                        DirectAssessmentCriterionScore(
+                            criterion="conciseness", score=4, explanation="explanation2"
+                        ),
+                        DirectAssessmentCriterionScore(
+                            criterion="relevance", score=5, explanation="explanation3"
+                        ),
+                    ],
+                ),
+                DirectAssessmentCriteriaScores(
+                    model="together://meta-llama/Llama-3.2-3B-Instruct-Turbo",
+                    criteria_scores=[
+                        DirectAssessmentCriterionScore(
+                            criterion="helpfulness", score=1, explanation="explanation1"
+                        ),
+                        DirectAssessmentCriterionScore(
+                            criterion="conciseness", score=2, explanation="explanation2"
+                        ),
+                        DirectAssessmentCriterionScore(
+                            criterion="relevance", score=3, explanation="explanation3"
+                        ),
+                    ],
+                ),
+            ]
+        )
+    else:
+        completion = client.beta.chat.completions.parse(
+            model="gpt-4o-mini",
+            messages=[
+                {
+                    "role": "system",
+                    "content": "Parse the judging responses into structured data.",
+                },
+                {"role": "user", "content": prompt},
+            ],
+            response_format=DirectAssessmentJudgingResponse,
+        )
+        return completion.choices[0].message.parsed
 
 
 def plot_criteria_scores(df):
@@ -364,6 +368,94 @@ def plot_criteria_scores(df):
     st.pyplot(plt.gcf())
 
 
+def plot_overall_scores(overall_scores_df):
+    # Calculate mean and standard deviation
+    summary = (
+        overall_scores_df.groupby("response_model")
+        .agg({"score": ["mean", "std"]})
+        .reset_index()
+    )
+    summary.columns = ["response_model", "mean_score", "std_score"]
+
+    # Add UI-friendly names
+    summary["ui_friendly_name"] = summary["response_model"].apply(get_ui_friendly_name)
+
+    # Sort the summary dataframe by mean_score in descending order
+    summary = summary.sort_values("mean_score", ascending=False)
+
+    # Create the plot
+    plt.figure(figsize=(8, 5))
+
+    # Plot bars with rainbow colors
+    ax = sns.barplot(
+        x="ui_friendly_name",
+        y="mean_score",
+        data=summary,
+        palette="prism",
+        capsize=0.1,
+    )
+
+    # Add error bars manually
+    x_coords = range(len(summary))
+    plt.errorbar(
+        x=x_coords,
+        y=summary["mean_score"],
+        yerr=summary["std_score"],
+        fmt="none",
+        c="black",
+        capsize=5,
+        zorder=10,  # Ensure error bars are on top
+    )
+
+    # Add text annotations
+    for i, row in summary.iterrows():
+        ax.text(
+            i,
+            row["mean_score"],
+            f"{row['mean_score']:.2f}",
+            ha="center",
+            va="bottom",
+            fontweight="bold",
+            color="black",
+            bbox=dict(facecolor="white", edgecolor="none", alpha=0.7, pad=0.5),
+        )
+
+    # Customize the plot
+    plt.xlabel("")
+    plt.ylabel("Overall Score")
+    plt.xticks(rotation=45, ha="right")
+    plt.tight_layout()
+
+    # Display the plot in Streamlit
+    st.pyplot(plt.gcf())
+
+
+def plot_per_judge_overall_scores(df):
+    # Find the overall score by finding the overall score for each judge, and then averaging
+    # over all judges.
+    grouped = df.groupby(["llm_judge_model"]).agg({"score": ["mean"]}).reset_index()
+    grouped.columns = ["llm_judge_model", "overall_score"]
+
+    # Create the horizontal bar plot
+    plt.figure(figsize=(10, 6))
+    ax = sns.barplot(
+        data=grouped,
+        y="llm_judge_model",
+        x="overall_score",
+        hue="llm_judge_model",
+        orient="h",
+    )
+
+    # Customize the plot
+    plt.title("Overall Scores by LLM Judge Model")
+    plt.xlabel("Overall Score")
+    plt.ylabel("LLM Judge Model")
+
+    # Adjust layout and display the plot
+    plt.tight_layout()
+    st.pyplot(plt)
+
+
 # Main Streamlit App
 def main():
     st.set_page_config(
@@ -395,7 +487,7 @@ def main():
 
     # App title and description
     st.title("Language Model Council Sandbox")
-    st.markdown("###### Invoke a council of LLMs to generate and judge each other.")
+    st.markdown("###### Invoke a council of LLMs to judge each other's responses.")
     st.markdown("###### [Paper](https://arxiv.org/abs/2406.08598)")
 
     # Authentication system
@@ -413,18 +505,30 @@ def main():
                     st.error("Invalid credentials")
 
     if st.session_state.authenticated:
-        st.success("Logged in successfully!")
+        # cols[1].success("Logged in successfully!")
+        st.markdown("#### LLM Council Member Selection")
 
         # Council and aggregator selection
         selected_models = llm_council_selector()
-        st.write("Selected Models:", selected_models)
+
+        # st.write("Selected Models:", selected_models)
+
         selected_aggregator = aggregator_selector()
 
+        # Initialize session state for collecting responses.
+        if "responses" not in st.session_state:
+            st.session_state.responses = {}
+        # if "aggregator_response" not in st.session_state:
+        # st.session_state.aggregator_response = {}
+
         # Prompt input
-        user_prompt = st.text_area("Enter your prompt:")
+        st.markdown("#### Enter your prompt")
+        _, center_column, _ = st.columns([3, 5, 3])
+        with center_column:
+            user_prompt = st.text_area(value="Say 'Hello World'", label="")
 
-        if st.button("Submit"):
-            st.write("Responses:")
+        if center_column.button("Submit", use_container_width=True):
+            st.markdown("#### Responses")
 
             response_columns = st.columns(3)
 
@@ -443,7 +547,7 @@ def main():
                         message_placeholder = st.empty()
                         stream = get_llm_response_stream(selected_model, user_prompt)
                         if stream:
-                            st.session_state[get_response_key(selected_model)] = (
+                            st.session_state["responses"][selected_model] = (
                                 message_placeholder.write_stream(stream)
                             )
 
@@ -456,25 +560,25 @@ def main():
                 st.code(aggregator_prompt)
 
             # Fetching and streaming response from the aggregator
-            st.write(
-                f"Mixture-of-Agents response from {get_ui_friendly_name(selected_aggregator)}"
-            )
+            st.write(f"Mixture-of-Agents ({get_ui_friendly_name(selected_aggregator)})")
             with st.chat_message(
                 selected_aggregator,
-                avatar=PROVIDER_TO_AVATAR_MAP[selected_aggregator],
+                avatar="img/council_icon.png",
             ):
                 message_placeholder = st.empty()
                 aggregator_stream = get_llm_response_stream(
                     selected_aggregator, aggregator_prompt
                 )
                 if aggregator_stream:
-                    message_placeholder.write_stream(aggregator_stream)
-                    st.session_state[
-                        get_aggregator_response_key(selected_aggregator)
-                    ] = message_placeholder.write_stream(aggregator_stream)
+                    st.session_state["responses"]["agg__" + selected_aggregator] = (
+                        message_placeholder.write_stream(aggregator_stream)
+                    )
+
+        # st.write("Responses (in session state):")
+        # st.write(st.session_state["responses"])
 
         # Judging.
-        st.markdown("#### Judging Configuration Form")
+        st.markdown("#### Judging Configuration")
 
         # Choose the type of assessment
         assessment_type = st.radio(
@@ -482,9 +586,48 @@ def main():
             options=["Direct Assessment", "Pairwise Comparison"],
         )
 
+        _, center_column, _ = st.columns([3, 5, 3])
+
         # Depending on the assessment type, render different forms
         if assessment_type == "Direct Assessment":
-            with st.expander("Direct Assessment Prompt"):
+
+            # 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),
@@ -495,10 +638,15 @@ def main():
             criteria_list = DEFAULT_DIRECT_ASSESSMENT_CRITERIA_LIST
 
             # Create DirectAssessment object when form is submitted
-            if st.button("Submit Direct Assessment"):
+            if center_column.button(
+                "Submit Direct Assessment", use_container_width=True
+            ):
 
                 # Submit direct asssessment.
-                responses_for_judging = get_response_mapping()
+                responses_for_judging = st.session_state["responses"]
+
+                # st.write("Responses for judging (in session state):")
+                # st.write(responses_for_judging)
 
                 response_judging_columns = st.columns(3)
 
@@ -515,11 +663,13 @@ def main():
                     ]
 
                     with st_column:
-                        if "aggregator_response" in response_model:
+                        if "agg__" in response_model:
                             judging_model_header = "Mixture-of-Agents Response"
                         else:
                             judging_model_header = get_ui_friendly_name(response_model)
                         st.write(f"Judging for {judging_model_header}")
+                        # st.write("Response being judged: ")
+                        # st.write(response)
                         judging_prompt = get_direct_assessment_prompt(
                             direct_assessment_prompt=direct_assessment_prompt,
                             user_prompt=user_prompt,
@@ -543,18 +693,27 @@ def main():
                                     judging_stream = get_llm_response_stream(
                                         judging_model, judging_prompt
                                     )
-                                    if judging_stream:
-                                        st.session_state[
-                                            get_direct_assessment_judging_key(
-                                                judging_model, response_model
-                                            )
-                                        ] = message_placeholder.write_stream(
-                                            judging_stream
-                                        )
+                                    # if judging_stream:
+                                    st.session_state[
+                                        "direct_assessment_judging_responses"
+                                    ][response_model][
+                                        judging_model
+                                    ] = message_placeholder.write_stream(
+                                        judging_stream
+                                    )
                         # When all of the judging is finished for the given response, get the actual
                         # values, parsed (use gpt-4o-mini for now) with json mode.
                         # TODO.
-                        judging_responses = get_direct_assessment_judging_responses()
+                        judging_responses = st.session_state[
+                            "direct_assessment_judging_responses"
+                        ][response_model]
+
+                        # st.write("Judging responses (in session state):")
+                        # st.write(judging_responses)
+
+                        if not judging_responses:
+                            st.error(f"No judging responses for {response_model}")
+                            quit()
                         parse_judging_response_prompt = (
                             get_parse_judging_response_for_direct_assessment_prompt(
                                 judging_responses,
@@ -562,45 +721,124 @@ def main():
                                 SEVEN_POINT_DIRECT_ASSESSMENT_OPTIONS,
                             )
                         )
+                        with st.expander("Parse Judging Response Prompt"):
+                            st.code(parse_judging_response_prompt)
                         # Issue the prompt to openai mini with structured outputs
                         parsed_judging_responses = parse_judging_responses(
-                            parse_judging_response_prompt
+                            parse_judging_response_prompt, judging_responses
                         )
 
-                        df = create_dataframe_for_direct_assessment_judging_response(
+                        st.session_state["direct_assessment_judging_df"][
+                            response_model
+                        ] = create_dataframe_for_direct_assessment_judging_response(
                             parsed_judging_responses
                         )
-                        st.write(df)
+                        st.write(
+                            st.session_state["direct_assessment_judging_df"][
+                                response_model
+                            ]
+                        )
 
-                        # Log the output using st.write() under an st.expander
-                        # with st.expander("Parsed Judging Responses", expanded=True):
-                        # st.write(parsed_judging_responses)
-                        plot_criteria_scores(df)
+                        plot_criteria_scores(
+                            st.session_state["direct_assessment_judging_df"][
+                                response_model
+                            ]
+                        )
 
-                        # TODO: Use parsed_judging_responses for further processing or display
+                        # Find the overall score by finding the overall score for each judge, and then averaging
+                        # over all judges.
+                        plot_per_judge_overall_scores(
+                            st.session_state["direct_assessment_judging_df"][
+                                response_model
+                            ]
+                        )
 
-        elif assessment_type == "Pairwise Comparison":
-            pairwise_comparison_prompt = st.text_area(
-                "Prompt for the Pairwise Comparison"
-            )
-            granularity = st.selectbox("Granularity", ["coarse", "fine", "super fine"])
-            ties_allowed = st.checkbox("Are ties allowed?")
-            position_swapping = st.checkbox("Enable position swapping?")
-            reference_model = st.text_input("Reference Model")
-
-            # Create PairwiseComparison object when form is submitted
-            if st.button("Submit Pairwise Comparison"):
-                pairwise_comparison_config = PairwiseComparison(
-                    type="pairwise_comparison",
-                    granularity=granularity,
-                    ties_allowed=ties_allowed,
-                    position_swapping=position_swapping,
-                    reference_model=reference_model,
-                    prompt=prompt,
+                        grouped = (
+                            st.session_state["direct_assessment_judging_df"][
+                                response_model
+                            ]
+                            .groupby(["llm_judge_model"])
+                            .agg({"score": ["mean"]})
+                            .reset_index()
+                        )
+                        grouped.columns = ["llm_judge_model", "overall_score"]
+
+                        # st.write(
+                        #     "Extracting overall scores from this grouped dataframe:"
+                        # )
+                        # st.write(grouped)
+
+                        # Save the overall scores to the session state.
+                        for record in grouped.to_dict(orient="records"):
+                            st.session_state["direct_assessment_overall_scores"][
+                                response_model
+                            ][record["llm_judge_model"]] = record["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}")
+
+                st.session_state["judging_status"] = "complete"
+
+            # Judging is complete.
+            st.write("#### Results")
+            # The session state now contains the overall scores for each response from each judge.
+            if st.session_state["judging_status"] == "complete":
+                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'])}"
                 )
-                st.success(f"Pairwise Comparison Created: {pairwise_comparison_config}")
-                # Submit pairwise comparison.
-                responses_for_judging = get_response_mapping()
+                # 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:
+                    st.dataframe(overall_scores_df)
+
+        elif assessment_type == "Pairwise Comparison":
+            pass
+            # pairwise_comparison_prompt = st.text_area(
+            #     "Prompt for the Pairwise Comparison"
+            # )
+            # granularity = st.selectbox("Granularity", ["coarse", "fine", "super fine"])
+            # ties_allowed = st.checkbox("Are ties allowed?")
+            # position_swapping = st.checkbox("Enable position swapping?")
+            # reference_model = st.text_input("Reference Model")
+
+            # # Create PairwiseComparison object when form is submitted
+            # if st.button("Submit Pairwise Comparison"):
+            #     pairwise_comparison_config = PairwiseComparison(
+            #         type="pairwise_comparison",
+            #         granularity=granularity,
+            #         ties_allowed=ties_allowed,
+            #         position_swapping=position_swapping,
+            #         reference_model=reference_model,
+            #         prompt=prompt,
+            #     )
+            #     st.success(f"Pairwise Comparison Created: {pairwise_comparison_config}")
+            #     # Submit pairwise comparison.
+            #     responses_for_judging = st.session_state["responses"]
 
     else:
         with cols[1]:

prompts.py

@@ -25,7 +25,10 @@ DEFAULT_AGGREGATOR_PROMPT = """We are trying to come up with the best response t
 Responses from other LLMs:
 {responses_from_other_llms}
 
-Please provide a response the combines the best aspects of the responses above."""
+Consider how you would combine the best aspects of the responses above into a single response.
+
+Directly provide your response to the user's query as if you were the original LLM. Do not mention that you are synthesizing the responses from other LLMs.
+"""
 
 
 DEFAULT_DIRECT_ASSESSMENT_PROMPT = """We are trying to assess the quality of a response to a user query.