File size: 8,396 Bytes
0bab47c
 
4f5d1cb
0bab47c
 
97fcd0c
9756440
 
4f5d1cb
0bab47c
 
 
da814b0
 
0bab47c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
9756440
0bab47c
 
 
1758388
0bab47c
1758388
 
 
0bab47c
9756440
1758388
9756440
1758388
0bab47c
 
9756440
 
 
0bab47c
9756440
0bab47c
4f5d1cb
9756440
4f5d1cb
 
 
 
1758388
0bab47c
da814b0
 
9756440
4f5d1cb
9756440
 
0bab47c
1758388
4f5d1cb
9756440
1758388
da814b0
4f5d1cb
 
e9811e1
4f5d1cb
da814b0
4f5d1cb
 
0bab47c
da814b0
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
0bab47c
da814b0
0bab47c
 
97fcd0c
 
 
0bab47c
 
da814b0
 
0bab47c
 
 
 
 
 
 
 
 
 
 
 
 
 
97fcd0c
 
e9811e1
97fcd0c
1758388
 
da814b0
 
1758388
 
 
da814b0
1758388
 
 
 
 
0bab47c
1758388
0bab47c
 
9756440
0bab47c
9756440
0bab47c
9756440
 
0bab47c
 
 
1758388
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
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


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,
        config: dict = {},
    ):
        super().__init__(app, workflow, ui_state, model_options, config)

        self.buzzer_state = gr.State(workflow.buzzer.model_dump())

        self.pipeline_state.change(
            lambda x: TossupPipelineState(**x).workflow.buzzer.model_dump(),
            inputs=[self.pipeline_state],
            outputs=[self.buzzer_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 = state.workflow.is_token_probs_supported(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 = state.workflow.get_answer_model(answer_var)
        return (
            state.model_dump(),
            gr.update(interactive=logprobs_supported, value=tokens_prob if logprobs_supported else 0.0),
            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 = pipeline_state.workflow.is_token_probs_supported()
            selected_model_name = pipeline_state.workflow.get_answer_model()
            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()),
        )

        # Updating the pipeline buzzer on user input changes in Buzzer panel.
        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],
        )

        # THIS WASN't NEEDED SINCE WE RERENDER THE OUTPUT PANEL ENTIRELY ON CHANGES
        # answer_dropdown = dropdowns["answer"]
        # if answer_dropdown is not None:
        #     gr.on(
        #         triggers=[self.buzzer_state.change],
        #         fn=self.update_prob_slider,
        #         inputs=[self.buzzer_state, answer_dropdown, self.prob_slider],
        #         outputs=[
        #             self.pipeline_state,
        #             self.prob_slider,
        #             self.buzzer_method_dropdown,
        #             self.buzzer_warning_display,
        #         ],
        #     )