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 display.formatting import tiny_styled_warning
from workflows.structs import Buzzer, TossupWorkflow

from .model_pipeline import PipelineInterface, PipelineState, PipelineUIState


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 TossupPipelineState(PipelineState):
    workflow: TossupWorkflow


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

    def update_buzzer(
        self,
        state: TossupPipelineState,
        confidence_threshold: float,
        method: str,
        tokens_prob: float | None,
    ):
        """Update the buzzer."""

        prob_threshold = float(tokens_prob) if tokens_prob and tokens_prob > 0 else None
        state.workflow.buzzer = state.workflow.buzzer.model_copy(
            update={
                "method": method,
                "confidence_threshold": confidence_threshold,
                "prob_threshold": prob_threshold,
            }
        )
        Buzzer.model_validate(state.workflow.buzzer)
        return state

    def update_prob_slider(self, state: TossupPipelineState, answer_var: str, tokens_prob: float | None):
        """Update the probability slider based on the answer variable."""
        if answer_var == UNSELECTED_VAR_NAME:
            return (
                state,
                gr.update(interactive=True),
                gr.update(value="AND", interactive=True),
                gr.update(visible=False),
            )
        step_id = answer_var.split(".")[0]
        model_name = state.workflow.steps[step_id].get_full_model_name()
        model_config = AVAILABLE_MODELS[model_name]
        is_model_with_logprobs = model_config.get("logprobs", False)
        buzzer = state.workflow.buzzer
        tokens_prob_threshold = tokens_prob if is_model_with_logprobs else None
        method = buzzer.method if is_model_with_logprobs else "AND"
        state = self.update_buzzer(
            state,
            confidence_threshold=buzzer.confidence_threshold,
            method=method,
            tokens_prob=tokens_prob_threshold,
        )
        return (
            state,
            gr.update(interactive=is_model_with_logprobs),
            gr.update(value=method, interactive=is_model_with_logprobs),
            gr.update(
                value=tiny_styled_warning(
                    f"{model_name} does not support `logprobs`. The probability slider will be disabled."
                ),
                visible=not is_model_with_logprobs,
            ),
        )

    def _render_output_panel(self, available_variables: list[str], pipeline_state: TossupPipelineState):
        dropdowns = {}
        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, 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.",
            )
            with gr.Row(elem_classes="control-panel"):
                self.confidence_slider = gr.Slider(
                    minimum=0.0,
                    maximum=1.0,
                    value=self.defaults.get("confidence_threshold", 0.85),
                    step=0.01,
                    label="Confidence",
                    elem_classes="slider-container",
                )
                self.buzzer_method_dropdown = gr.Dropdown(
                    choices=["AND", "OR"],
                    value=self.defaults.get("buzzer_method", "AND"),
                    label="Method",
                    interactive=True,
                    min_width=80,
                    scale=0,
                )
                self.prob_slider = gr.Slider(
                    value=self.defaults.get("logits_prob", 0.0),
                    label="Probability",
                    minimum=0.0,
                    maximum=1.0,
                    step=0.001,
                    elem_classes="slider-container",
                )
            self.buzzer_warning_display = gr.HTML(visible=False)

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

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

        gr.on(
            triggers=[
                self.confidence_slider.input,
                self.buzzer_method_dropdown.input,
                self.prob_slider.input,
            ],
            fn=self.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,
                ],
            )