src/components/model_pipeline/tossup_pipeline.py

import gradio as gr
import numpy as np
from loguru import logger

from app_configs import AVAILABLE_MODELS, UNSELECTED_VAR_NAME
from components import commons
from components.structs import TossupPipelineState
from components.typed_dicts import TossupPipelineStateDict
from display.formatting import tiny_styled_warning
from workflows.structs import Buzzer, TossupWorkflow

from .model_pipeline import PipelineInterface, PipelineState, PipelineUIState
from .state_manager import PipelineStateManager, TossupPipelineStateManager


def get_probs_model_name(workflow: TossupWorkflow, answer_var: str | None = None) -> str | None:
    if answer_var is None:
        answer_var = workflow.outputs["answer"]
    if answer_var is None or answer_var == UNSELECTED_VAR_NAME:
        return None
    step_id = answer_var.split(".")[0]
    return workflow.steps[step_id].get_full_model_name()


def is_logprobs_supported(workflow: TossupWorkflow, answer_var: str | None = None) -> bool:
    model_name = get_probs_model_name(workflow, answer_var)
    if model_name is None:
        return True
    return AVAILABLE_MODELS[model_name].get("logprobs", False)


def toggleable_slider(
    value, minimum, maximum, step, toggle_value=False, label=None, info=None, min_width=200, scale=1
):
    with gr.Column(elem_classes="toggleable", min_width=min_width, scale=scale):
        show_label = label is not None
        checkbox = gr.Checkbox(label=label, value=toggle_value, container=False, info=info, show_label=show_label)
        slider = gr.Slider(
            minimum=minimum,
            maximum=maximum,
            value=value,
            step=step,
            label="",
            interactive=True,
            show_label=False,
            container=False,
        )
        checkbox.change(fn=lambda x: gr.update(interactive=x), inputs=[checkbox], outputs=[slider])
    return checkbox, slider


class TossupPipelineInterface(PipelineInterface):
    def __init__(
        self,
        app: gr.Blocks,
        workflow: TossupWorkflow,
        ui_state: PipelineUIState | None = None,
        model_options: list[str] = None,
        simple: bool = False,
        defaults: dict = {},
    ):
        super().__init__(app, workflow, ui_state, model_options, simple)
        self.defaults = defaults

        self.pipeline_state.change(
            lambda x: logger.debug(
                f"Pipeline state changed. Type: {type(x)}. Has buzzer info: {x['workflow']['buzzer'] if isinstance(x, dict) else 'N/A'}"
            ),
            inputs=[self.pipeline_state],
        )

    def update_prob_slider(
        self, state_dict: TossupPipelineStateDict, answer_var: str, tokens_prob: float | None
    ) -> tuple[TossupPipelineStateDict, dict, dict, dict]:
        """Update the probability slider based on the answer variable."""
        state = TossupPipelineState(**state_dict)
        if answer_var == UNSELECTED_VAR_NAME:
            return (
                state.model_dump(),
                gr.update(interactive=True),
                gr.update(value="AND", interactive=True),
                gr.update(visible=False),
            )
        logprobs_supported = is_logprobs_supported(state.workflow, answer_var)
        buzzer = state.workflow.buzzer
        tokens_prob_threshold = tokens_prob if logprobs_supported else None
        method = buzzer.method if logprobs_supported else "AND"
        state.workflow.buzzer = Buzzer(
            method=method,
            confidence_threshold=buzzer.confidence_threshold,
            prob_threshold=tokens_prob_threshold,
        )
        model_name = get_probs_model_name(state.workflow, answer_var)
        return (
            state.model_dump(),
            gr.update(interactive=logprobs_supported),
            gr.update(value=method, interactive=logprobs_supported),
            gr.update(
                value=tiny_styled_warning(
                    f"<code>'{model_name}'</code> does not support <code>'logprobs'</code>. The probability slider will be disabled."
                ),
                visible=not logprobs_supported,
            ),
        )

    def _render_buzzer_panel(
        self, buzzer: Buzzer, prob_slider_supported: bool, selected_model_name: str | None = None
    ):
        with gr.Row(elem_classes="control-panel"):
            self.confidence_slider = gr.Slider(
                minimum=0.0,
                maximum=1.0,
                value=buzzer.confidence_threshold,
                step=0.01,
                label="Confidence",
                elem_classes="slider-container",
                show_reset_button=False,
            )
            value = buzzer.method if prob_slider_supported else "AND"
            self.buzzer_method_dropdown = gr.Dropdown(
                choices=["AND", "OR"],
                value=value,
                label="Method",
                interactive=prob_slider_supported,
                min_width=80,
                scale=0,
            )
            self.prob_slider = gr.Slider(
                value=buzzer.prob_threshold or 0.0,
                interactive=prob_slider_supported,
                label="Probability",
                minimum=0.0,
                maximum=1.0,
                step=0.001,
                elem_classes="slider-container",
                show_reset_button=False,
            )
        display_html = ""
        if selected_model_name is not None:
            display_html = tiny_styled_warning(
                f"<code>{selected_model_name}</code> does not support <code>logprobs</code>. The probability slider will be disabled."
            )
        self.buzzer_warning_display = gr.HTML(display_html, visible=not prob_slider_supported)

    def _render_output_panel(self, pipeline_state: TossupPipelineState):
        dropdowns = {}
        available_variables = pipeline_state.workflow.get_available_variables()
        variable_options = [UNSELECTED_VAR_NAME] + [v for v in available_variables if v not in self.input_variables]
        with gr.Column(elem_classes="step-accordion control-panel"):
            commons.get_panel_header(
                header="Final output variables mapping:",
            )
            with gr.Row(elem_classes="output-fields-row"):
                for output_field in self.required_output_variables:
                    value = pipeline_state.workflow.outputs.get(output_field)
                    value = value or UNSELECTED_VAR_NAME
                    dropdown = gr.Dropdown(
                        label=output_field,
                        value=value,
                        choices=variable_options,
                        interactive=True,
                        elem_classes="output-field-variable",
                        # show_label=False,
                    )
                    dropdown.change(
                        self.sm.update_output_variables,
                        inputs=[self.pipeline_state, gr.State(output_field), dropdown],
                        outputs=[self.pipeline_state],
                    )
                    dropdowns[output_field] = dropdown
            commons.get_panel_header(
                header="Buzzer settings:",
                subheader="Set your thresholds for confidence and output tokens probability (computed using <code>logprobs</code>).",
            )
            logprobs_supported = is_logprobs_supported(pipeline_state.workflow)
            selected_model_name = get_probs_model_name(pipeline_state.workflow)
            self._render_buzzer_panel(pipeline_state.workflow.buzzer, logprobs_supported, selected_model_name)

        # def update_choices(available_variables: list[str]):
        #     """Update the choices for the dropdowns"""
        #     return [gr.update(choices=available_variables, value=None, selected=None) for _ in dropdowns.values()]

        # self.variables_state.change(
        #     update_choices,
        #     inputs=[self.variables_state],
        #     outputs=list(dropdowns.values()),
        # )

        gr.on(
            triggers=[
                self.confidence_slider.release,
                self.buzzer_method_dropdown.input,
                self.prob_slider.release,
            ],
            fn=self.sm.update_buzzer,
            inputs=[self.pipeline_state, self.confidence_slider, self.buzzer_method_dropdown, self.prob_slider],
            outputs=[self.pipeline_state],
        )

        # TODO: Do Add model step change triggers as well. (Model name change triggers)
        answer_dropdown = dropdowns["answer"]
        if answer_dropdown is not None:
            gr.on(
                triggers=[answer_dropdown.input, self.model_selection_state.change],
                fn=self.update_prob_slider,
                inputs=[self.pipeline_state, answer_dropdown, self.prob_slider],
                outputs=[
                    self.pipeline_state,
                    self.prob_slider,
                    self.buzzer_method_dropdown,
                    self.buzzer_warning_display,
                ],
            )