minor interface refactor

src/components/quizbowl/bonus.py

@@ -2,6 +2,7 @@ import json
 from typing import Any
 
 import gradio as gr
+import numpy as np
 import pandas as pd
 from datasets import Dataset
 from loguru import logger
@@ -36,7 +37,7 @@ def process_bonus_results(results: list[dict]) -> pd.DataFrame:
     )
 
 
-def initialize_eval_interface(example: dict, model_outputs: list[dict]):
+def initialize_eval_interface(example: dict, model_outputs: list[dict], input_vars: list[str]):
     """Initialize the interface with example text."""
     try:
         html_content = create_bonus_html(example["leadin"], example["parts"])
@@ -45,12 +46,20 @@ def initialize_eval_interface(example: dict, model_outputs: list[dict]):
         plot_data = create_bonus_confidence_plot(example["parts"], model_outputs)
 
         # Store state
-        state = json.dumps({"parts": example["parts"], "outputs": model_outputs})
+        state = {"parts": example["parts"], "outputs": model_outputs}
 
-        return html_content, plot_data, state
+        # Preparing step outputs for the model
+        step_outputs = {}
+        for i, output in enumerate(model_outputs):
+            key = f"part {i + 1}"
+            step_outputs[key] = {k: v for k, v in output["step_outputs"].items() if k not in input_vars}
+            if output["logprob"] is not None:
+                step_outputs[key]["output_probability"] = float(np.exp(output["logprob"]))
+
+        return html_content, plot_data, state, step_outputs
     except Exception as e:
         logger.exception(f"Error initializing interface: {e.args}")
-        return f"<div>Error initializing interface: {str(e)}</div>", pd.DataFrame(), "{}"
+        return f"<div>Error initializing interface: {str(e)}</div>", pd.DataFrame(), {}, {}
 
 
 class BonusInterface:
@@ -64,7 +73,7 @@ class BonusInterface:
         self.model_options = model_options
         self.app = app
         self.defaults = defaults
-        self.output_state = gr.State(value="{}")
+        self.output_state = gr.State(value={})
         self.render()
 
     # ------------------------------------- LOAD PIPELINE STATE FROM BROWSER STATE -------------------------------------
@@ -75,10 +84,10 @@ class BonusInterface:
             state_dict = browser_state["bonus"].get("pipeline_state", {})
             pipeline_state = PipelineState.model_validate(state_dict)
             pipeline_state_dict = pipeline_state.model_dump()
-            output_state = browser_state["bonus"].get("output_state", "{}")
+            output_state = browser_state["bonus"].get("output_state", {})
         except Exception as e:
             logger.warning(f"Error loading presaved pipeline state: {e}")
-            output_state = "{}"
+            output_state = {}
             workflow = self.defaults["init_workflow"]
             pipeline_state_dict = PipelineState.from_workflow(workflow).model_dump()
         return browser_state, not pipeline_change, pipeline_state_dict, output_state
@@ -228,9 +237,10 @@ class BonusInterface:
             outputs = self.get_agent_outputs(example, pipeline_state)
 
             # Process results and prepare visualization data
-            html_content, plot_data, output_state = initialize_eval_interface(example, outputs)
+            html_content, plot_data, output_state, step_outputs = initialize_eval_interface(
+                example, outputs, pipeline_state.workflow.inputs
+            )
             df = process_bonus_results(outputs)
-            step_outputs = [output["step_outputs"] for output in outputs]
 
             return (
                 html_content,

src/components/quizbowl/tossup.py

@@ -37,12 +37,12 @@ class ScoredTossupResult(TossupResult):
     token_position: int  # Position in the question where prediction was made
 
 
-def add_model_scores(model_outputs: list[dict], clean_answers: list[str], run_indices: list[int]) -> list[dict]:
+def add_model_scores(run_outputs: list[dict], clean_answers: list[str], run_indices: list[int]) -> list[dict]:
     """Add model scores to the model outputs."""
-    for output, run_idx in zip(model_outputs, run_indices):
+    for output in run_outputs:
         output["score"] = evaluate_prediction(output["answer"], clean_answers)
-        output["token_position"] = run_idx + 1
-    return model_outputs
+        output["token_position"] = run_indices[output["position"] - 1]
+    return run_outputs
 
 
 def prepare_buzz_evals(
@@ -61,7 +61,11 @@ def prepare_buzz_evals(
 
 
 def initialize_eval_interface(
-    example: dict, model_outputs: list[dict], confidence_threshold: float, prob_threshold: float | None = None
+    example: dict,
+    run_outputs: list[dict],
+    input_vars: list,
+    confidence_threshold: float,
+    prob_threshold: float | None = None,
 ):
     """Initialize the interface with example text."""
     try:
@@ -69,7 +73,7 @@ def initialize_eval_interface(
         run_indices = example["run_indices"]
         answer = example["answer_primary"]
         clean_answers = example["clean_answers"]
-        eval_points = prepare_buzz_evals(run_indices, model_outputs)
+        eval_points = [(o["token_position"], o) for o in run_outputs]
 
         if not tokens:
             return "<div>No tokens found in the provided text.</div>", pd.DataFrame(), "{}"
@@ -77,12 +81,21 @@ def initialize_eval_interface(
         plot_data = create_tossup_confidence_pyplot(tokens, eval_points, confidence_threshold, prob_threshold)
 
         # Store tokens, values, and buzzes as JSON for later use
-        state = json.dumps({"tokens": tokens, "values": eval_points})
-
-        return html_content, plot_data, state
+        state = {"tokens": tokens, "values": eval_points}
+
+        # Preparing step outputs for the model
+        step_outputs = {}
+        for output in run_outputs:
+            tok_pos = output["token_position"]
+            key = "{pos}:{token}".format(pos=tok_pos + 1, token=tokens[tok_pos])
+            step_outputs[key] = {k: v for k, v in output["step_outputs"].items() if k not in input_vars}
+            if output["logprob"] is not None:
+                step_outputs[key]["output_probability"] = float(np.exp(output["logprob"]))
+
+        return html_content, plot_data, state, step_outputs
     except Exception as e:
         logger.exception(f"Error initializing interface: {e.args}")
-        return f"<div>Error initializing interface: {str(e)}</div>", pd.DataFrame(), "{}"
+        return f"<div>Error initializing interface: {str(e)}</div>", pd.DataFrame(), "{}", {}
 
 
 def process_tossup_results(results: list[dict]) -> pd.DataFrame:
@@ -119,7 +132,7 @@ class TossupInterface:
         self.model_options = model_options
         self.app = app
         self.defaults = defaults
-        self.output_state = gr.State(value="{}")
+        self.output_state = gr.State(value={})
         self.render()
 
     # ------------------------------------- LOAD PIPELINE STATE FROM BROWSER STATE -------------------------------------
@@ -130,10 +143,10 @@ class TossupInterface:
             state_dict = browser_state["tossup"].get("pipeline_state", {})
             pipeline_state = TossupPipelineState.model_validate(state_dict)
             pipeline_state_dict = pipeline_state.model_dump()
-            output_state = browser_state["tossup"].get("output_state", "{}")
+            output_state = browser_state["tossup"].get("output_state", {})
         except Exception as e:
             logger.warning(f"Error loading presaved pipeline state: {e}")
-            output_state = "{}"
+            output_state = {}
             workflow = self.defaults["init_workflow"]
             pipeline_state_dict = TossupPipelineState.from_workflow(workflow).model_dump()
         return browser_state, not pipeline_change, pipeline_state_dict, output_state
@@ -282,20 +295,10 @@ class TossupInterface:
             # Process results and prepare visualization data
             confidence_threshold = workflow.buzzer.confidence_threshold
             prob_threshold = workflow.buzzer.prob_threshold
-            tokens_html, plot_data, output_state = initialize_eval_interface(
-                example, outputs, confidence_threshold, prob_threshold
+            tokens_html, plot_data, output_state, step_outputs = initialize_eval_interface(
+                example, outputs, workflow.inputs, confidence_threshold, prob_threshold
             )
             df = process_tossup_results(outputs)
-            tokens = example["question"].split()
-            step_outputs = {}
-            for output in outputs:
-                pos = output["token_position"]
-                token = tokens[pos - 1]
-                key = f"{pos}:{token}"
-                step_outputs[key] = {k: v for k, v in output["step_outputs"].items() if k not in workflow.inputs}
-                if output["logprob"] is not None:
-                    step_outputs[key]["logprob"] = output["logprob"]
-                    step_outputs[key]["prob"] = float(np.exp(output["logprob"]))
 
             return (
                 tokens_html,