ui.py

"""Gradio UI for bill vs non-bill v3_teacher (teacher / DeBERTa LoRA)."""

from __future__ import annotations

import os
from typing import Any, Tuple

import gradio as gr

from csv_inference import run_csv_inference, validate_csv_path
from inference import classify_one

DEFAULT_BILL_TEXT = (
    "Your bill of Rs. 1,250.00 for account ****42 is due by 25-Apr-2026. "
    "Pay via MyProvider app or https://pay.example.com to avoid late fees."
)


def _file_path(file_obj: object) -> str | None:
    if file_obj is None:
        return None
    if isinstance(file_obj, str):
        return file_obj if os.path.isfile(file_obj) else None
    path = getattr(file_obj, "name", None)
    if path and os.path.isfile(path):
        return path
    return None


def on_csv_change(file_obj: object) -> Tuple[Any, Any, Any]:
    path = _file_path(file_obj)
    if not path:
        return (
            gr.update(interactive=False),
            gr.update(value="Upload a CSV file to begin."),
            gr.update(value=None, interactive=False),
        )
    ok, msg = validate_csv_path(path)
    if not ok:
        print("CSV validation:", msg)
        return (
            gr.update(interactive=False),
            gr.update(value=f"Validation failed: {msg}"),
            gr.update(value=None, interactive=False),
        )
    return (
        gr.update(interactive=True),
        gr.update(value="File looks good. Click Process CSV."),
        gr.update(value=None, interactive=False),
    )


def run_text_classify(text: str) -> str:
    return classify_one(text)


def run_csv_start(file_obj: object, progress: gr.Progress = gr.Progress()) -> Tuple[Any, Any]:
    path = _file_path(file_obj)
    if not path:
        return gr.update(interactive=False, value=None), gr.update(value="No file uploaded.")
    try:
        def cb(done: int, total: int) -> None:
            if total > 0:
                progress((done, total), desc=f"{done} / {total}", total=total)

        out_path = run_csv_inference(path, progress_cb=cb)
        return (
            gr.update(interactive=True, value=out_path),
            gr.update(value="Done. Download your CSV below."),
        )
    except Exception as exc:
        print("run_csv_start:", exc)
        return (
            gr.update(interactive=False, value=None),
            gr.update(value=f"Error: {exc}"),
        )


def build_demo() -> gr.Blocks:
    with gr.Blocks(title="Bill vs non-bill v3_teacher (teacher)") as demo:
        gr.Markdown("## Bill vs non-bill v3_teacher (teacher)")

        with gr.Tabs():
            with gr.Tab("Text"):
                inp = gr.Textbox(
                    label="Message",
                    lines=6,
                    value=DEFAULT_BILL_TEXT,
                )
                go = gr.Button("Classify", variant="primary")
                out = gr.Textbox(label="Result JSON", lines=14)

            with gr.Tab("CSV"):
                csv_file = gr.File(label="CSV file", file_types=[".csv"], interactive=True)
                csv_status = gr.Textbox(
                    label="Status",
                    value="Upload a CSV with a `text` column.",
                    interactive=False,
                )
                start_btn = gr.Button("Process CSV", interactive=False, variant="primary")
                dl = gr.DownloadButton(label="Download result CSV", interactive=False)

        go.click(run_text_classify, inputs=[inp], outputs=[out])
        csv_file.change(on_csv_change, inputs=[csv_file], outputs=[start_btn, csv_status, dl])
        start_btn.click(run_csv_start, inputs=[csv_file], outputs=[dl, csv_status])

    return demo