app/ui.py

from pathlib import Path

import gradio as gr

from .config import REFRESH_SECONDS, display
from .handlers import EXAMPLES, classify, evaluate, pick_revision, refresh, warm
from .hub import model_repos, newest, revisions
from .text import (
    CLASSIFY,
    EVALUATE,
    HEADER,
    PLACEHOLDER,
    STATUS_LOADING,
    TRANSCRIPT_INFO,
    pushed_at,
    status_line,
)

CMD_ENTER_JS = (Path(__file__).parent / "cmd_enter.js").read_text(encoding="utf-8")


def selectors() -> tuple[gr.Dropdown, gr.Dropdown, gr.Button, gr.Markdown, gr.Markdown]:
    pushes = model_repos()
    repo = newest(pushes)
    revs = revisions(repo) if repo else []
    with gr.Row():
        model = gr.Dropdown(label="Model", choices=list(pushes), value=repo, scale=3)
        revision = gr.Dropdown(
            label="Revision",
            choices=revs,
            value=revs[0][1] if revs else None,
            scale=4,
        )
        # Beside Revision, because that is what it reports on. Never empty, so
        # Gradio's progress animation has something to sit on.
        status = gr.Markdown(status_line(STATUS_LOADING), scale=0, min_width=120)
        reload_button = gr.Button("↻", scale=0, min_width=48)
    pushed = gr.Markdown(pushed_at(pushes[repo] if repo else None))
    return model, revision, reload_button, pushed, status


def classify_tab() -> tuple[gr.Textbox, gr.Button, list]:
    with gr.Row():
        with gr.Column(scale=3):
            transcript = gr.Textbox(
                label="Caller's turns",
                info=TRANSCRIPT_INFO,
                placeholder=PLACEHOLDER,
                lines=7,
                max_lines=14,
            )
            run = gr.Button(
                CLASSIFY, variant="primary", elem_id="run-classify", interactive=False
            )
            gr.Examples(
                examples=[[row["text"]] for row in EXAMPLES],
                example_labels=[f"{display(row['gold'])} — {row['id']}" for row in EXAMPLES],
                label="Examples from the test split",
                inputs=transcript,
            )

        with gr.Column(scale=2):
            prediction = gr.Label(label="Prediction", num_top_classes=2)
            reads = gr.Textbox(
                label="What the model reads",
                lines=3,
                interactive=False,
                buttons=["copy"],
            )
            latency = gr.Markdown()

    return transcript, run, [prediction, reads, latency]


def evaluation_tab() -> tuple[gr.Button, list, gr.Dataframe]:
    run = gr.Button(EVALUATE, variant="primary", interactive=False)
    score = gr.Markdown()
    matrix = gr.Dataframe(
        label="Confusion matrix",
        headers=["", f"predicted {display('risk').lower()}", f"predicted {display('no_risk').lower()}"],
        interactive=False,
    )
    cases = gr.Dataframe(
        label="Cases (errors first)",
        headers=["", "id", "expected", "predicted", "confidence", "text"],
        datatype=["str", "str", "str", "str", "number", "str"],
        wrap=True,
        interactive=False,
    )
    return run, [score, matrix, cases], cases


def build() -> gr.Blocks:
    with gr.Blocks(title="Breakdown risk") as demo:
        gr.Markdown(HEADER)

        model, revision, refresh_button, pushed, status = selectors()
        selection = [model, revision]

        with gr.Tab("Classify"):
            transcript, run, results = classify_tab()

        with gr.Tab("Evaluation"):
            evaluate_button, report, progress_target = evaluation_tab()

        # Buttons start dead and the model is fetched up front, so a click is never
        # a download. Every pick warms the same way: "multiple" and no concurrency
        # limit because Gradio otherwise drops a pick made while one is pending
        # ("once") or queues it behind that download; warm() itself decides which
        # pick still owns the buttons. show_progress_on puts Gradio's own animation
        # on the little indicator beside Revision.
        warming = dict(
            outputs=[run, evaluate_button, status],
            show_progress="full",
            show_progress_on=[status],
            trigger_mode="multiple",
            concurrency_limit=None,
        )
        demo.load(warm, selection, **warming)
        # .input() is user-only: the revision update that picking a model triggers
        # must not warm a second time.
        revision.input(warm, selection, **warming)

        timer = gr.Timer(REFRESH_SECONDS)
        timer.tick(refresh, selection, [*selection, pushed], show_progress="hidden")
        refresh_button.click(refresh, selection, [*selection, pushed])
        model.change(pick_revision, model, [revision, pushed]).then(warm, selection, **warming)
        gr.on(
            [run.click, transcript.submit],
            classify,
            [*selection, transcript],
            results,
            api_name="classify",
        )
        evaluate_button.click(
            evaluate,
            selection,
            report,
            show_progress="full",
            show_progress_on=progress_target,
            api_name="evaluate",
        )

    return demo