File size: 10,315 Bytes
d9af7ee
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
"""Gradio entry-point wiring the patch-view service layer."""

from __future__ import annotations

import os
import socket
from typing import Dict, List

import gradio as gr

from .constraints import ConstraintDefinition
from .services import PatchService


def create_app(service: PatchService | None = None) -> gr.Blocks:
    """Create and configure the Gradio application."""
    patch_service = service or PatchService.from_data_root()
    audio_choices = patch_service.audio_ids()

    with gr.Blocks(title="BLT Patching Visualization", theme=gr.themes.Default()) as app:
        gr.Markdown("# BLT Patching Visualization Explorer")
        gr.Markdown("Explore entropy-based patching strategies for Byte Latent Transformers")

        patches_state = gr.State([])

        with gr.Row():
            with gr.Column(scale=1):
                gr.Markdown("## Audio Selection")

                audio_dropdown = gr.Dropdown(
                    choices=audio_choices,
                    value=audio_choices[0] if audio_choices else None,
                    label="Audio File",
                    interactive=True,
                )
                file_info = gr.Markdown()

                gr.Markdown("## Patching Parameters")

                metric_radio = gr.Radio(
                    choices=["Entropy", "Perplexity"],
                    value="Entropy",
                    label="Metric",
                    info="Underlying metric for patching decisions",
                )

                gr.Markdown("**Constraints** (select one or more)")
                gr.Markdown("_Patch boundary created when ANY constraint is satisfied (OR logic)_")

                constraint_defs: List[ConstraintDefinition] = list(
                    patch_service.constraint_definitions()
                )
                constraint_choices = [(definition.label, definition.name) for definition in constraint_defs]
                constraint_selector = gr.CheckboxGroup(
                    choices=constraint_choices,
                    value=["global"],
                    label="Active Constraints",
                    info="Select multiple constraints to combine with OR logic",
                )

                guard_config = patch_service.guard_slider_config("entropy")
                min_metric_slider = gr.Slider(
                    value=guard_config.default,
                    minimum=guard_config.minimum,
                    maximum=guard_config.maximum,
                    step=guard_config.step,
                    label=guard_config.label,
                    info=guard_config.description,
                )
                gr.Markdown("_Guard constraint uses AND logic with all selected constraints_")

                threshold_inputs: Dict[str, gr.Slider] = {}
                for definition in constraint_defs:
                    range_cfg = definition.range_for_metric("entropy")
                    threshold_inputs[definition.name] = gr.Slider(
                        value=range_cfg.default,
                        minimum=range_cfg.minimum,
                        maximum=range_cfg.maximum,
                        step=range_cfg.step,
                        label=f"{definition.label} Threshold",
                        info=definition.description,
                        visible=(definition.name == "global"),
                    )
                threshold_names = list(threshold_inputs.keys())

                apply_btn = gr.Button("Apply Patching", variant="primary", size="lg")

            with gr.Column(scale=1):
                stats_output = gr.Markdown("Click 'Apply Patching' to see statistics")

        gr.Markdown("## Visualization")
        plot_output = gr.Plot(label="Patch Visualization")

        gr.Markdown("## Audio Playback & Comparison")
        with gr.Tab("Original vs Reconstructed"):
            gr.Markdown(
                "**Compare original audio with WavTokenizer reconstruction** "
                "(listen for compression artifacts)"
            )
            with gr.Row():
                original_audio = gr.Audio(label="Original Audio", type="numpy")
                reconstructed_audio = gr.Audio(label="WavTokenizer Reconstructed", type="numpy")

        with gr.Tab("Play Until Boundary"):
            gr.Markdown(
                "**Play audio from a token position until the next patch boundary** "
                "(hear what's in a single patch)"
            )

            with gr.Row():
                current_token_input = gr.Number(
                    value=0,
                    label="Token Position",
                    precision=0,
                    info="Enter token index (0 to num_tokens-1)",
                )
                play_boundary_btn = gr.Button("▶ Play Until Next Boundary", variant="secondary")

            patch_info_output = gr.Markdown()
            boundary_audio = gr.Audio(label="Patch Audio", type="numpy", autoplay=True)

        # ------------------------------------------------------------------
        # Callbacks
        # ------------------------------------------------------------------
        def _update_file_info(file_id: str):
            return patch_service.file_info_markdown(file_id)

        def _update_threshold_visibility(selected: List[str]):
            return {
                threshold_inputs[name]: gr.update(visible=(name in selected))
                for name in threshold_inputs
            }

        def _update_threshold_ranges(metric_name: str):
            metric_defaults = patch_service.threshold_defaults(metric_name)
            guard = patch_service.guard_slider_config(metric_name)
            updates = {
                min_metric_slider: gr.update(
                    value=guard.default,
                    minimum=guard.minimum,
                    maximum=guard.maximum,
                    step=guard.step,
                    label=guard.label,
                    info=guard.description,
                )
            }
            for definition in constraint_defs:
                params = definition.range_for_metric(metric_name)
                updates[threshold_inputs[definition.name]] = gr.update(
                    value=metric_defaults[definition.name],
                    minimum=params.minimum,
                    maximum=params.maximum,
                    step=params.step,
                )
            return updates

        def _handle_patching(file_id, metric_name, selected_constraints, min_threshold, *slider_values):
            if not file_id:
                return None, "No file selected", None, None, []

            thresholds = {name: float(slider_values[idx]) for idx, name in enumerate(threshold_names)}

            try:
                return patch_service.run_patching(
                    file_id=file_id,
                    metric_name=metric_name,
                    selected_constraints=selected_constraints,
                    guard_threshold=min_threshold,
                    thresholds=thresholds,
                )
            except ValueError as exc:
                return None, f"⚠️ {exc}", None, None, []

        def _handle_play(file_id, current_token, patches):
            if not file_id:
                return None, "No file selected", 0
            if not patches:
                return None, "No patches computed yet - click 'Apply Patching' first", current_token
            try:
                return patch_service.play_until_boundary(file_id, int(current_token), patches)
            except ValueError as exc:
                return None, f"⚠️ {exc}", current_token

        audio_dropdown.change(fn=_update_file_info, inputs=[audio_dropdown], outputs=[file_info])

        constraint_selector.change(
            fn=_update_threshold_visibility,
            inputs=[constraint_selector],
            outputs=list(threshold_inputs.values()),
        )

        metric_radio.change(
            fn=_update_threshold_ranges,
            inputs=[metric_radio],
            outputs=[min_metric_slider, *threshold_inputs.values()],
        )

        apply_btn.click(
            fn=_handle_patching,
            inputs=[
                audio_dropdown,
                metric_radio,
                constraint_selector,
                min_metric_slider,
                *threshold_inputs.values(),
            ],
            outputs=[plot_output, stats_output, original_audio, reconstructed_audio, patches_state],
        )

        play_boundary_btn.click(
            fn=_handle_play,
            inputs=[audio_dropdown, current_token_input, patches_state],
            outputs=[boundary_audio, patch_info_output, current_token_input],
        )

        app.load(fn=_update_file_info, inputs=[audio_dropdown], outputs=[file_info])

    return app


def _find_open_port(preferred: int, attempts: int = 50) -> int:
    """Find an available TCP port starting from a preferred value."""
    if preferred <= 0:
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
            try:
                sock.bind(("127.0.0.1", 0))
                return sock.getsockname()[1]
            except OSError:
                return 7860

    port = preferred
    for _ in range(attempts):
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
            try:
                sock.bind(("127.0.0.1", port))
            except OSError:
                port += 1
                continue
            return port
    raise OSError(f"Could not find open port starting at {preferred}")


def main() -> None:
    """CLI entry point for running the Gradio app."""
    app = create_app()

    port_env = os.getenv("GRADIO_SERVER_PORT")
    try:
        requested_port = int(port_env) if port_env is not None else 7860
    except ValueError:
        requested_port = 7860

    try:
        port = _find_open_port(requested_port)
    except OSError as exc:
        raise SystemExit(f"Failed to locate an open port near {requested_port}: {exc}") from exc

    launch_kwargs = dict(server_name="0.0.0.0", share=True, debug=True, server_port=port)
    print(f"Open in browser: http://localhost:{port}")
    app.launch(**launch_kwargs)


# Gradio Spaces expects a variable named `demo`
demo = create_app()


if __name__ == "__main__":
    main()