app.py

#!/usr/bin/env python3
"""Gradio demo: a customer banking query in, the predicted intent out.

Runs **both** trained models side by side — LoRA and full fine-tuning — so the
project's claim is visible rather than asserted: near-identical predictions, from
models that differ by 321x in trainable parameters.

    python app.py                     # http://127.0.0.1:7860
    python app.py --share             # temporary public link
    python app.py --lora-only         # skip the 253 MB full model

If ``checkpoints/full.pt`` is missing the app degrades to a single-model view, so
a Hugging Face Space can ship only the 817 KB adapter (see DEPLOY.md).

WHY THIS DEPLOYS CHEAPLY
------------------------
The LoRA checkpoint holds only the adapter matrices and the classification head
(~817 KB). The frozen encoder is byte-identical to the public
``distilbert-base-uncased``, so the Space downloads that from the Hub at startup
rather than shipping a 253 MB copy of it. Staying under the 10 MB threshold means
no Git LFS and a fast push.
"""

from __future__ import annotations

import argparse
import json
import logging
import sys
import time
from pathlib import Path
from typing import Sequence

import gradio as gr
import torch
from transformers import AutoTokenizer

from models.classifier import TextClassifier

log = logging.getLogger("app")

DEFAULT_LORA = Path("checkpoints/lora.pt")
DEFAULT_FULL = Path("checkpoints/full.pt")
RESULTS_DIR = Path("results")
MAX_LENGTH = 64
TOP_K = 5

DISCLAIMER = """
<b>What this is.</b> A demonstration of LoRA implemented from scratch — no
<code>peft</code> — on the public
<a href="https://huggingface.co/datasets/mteb/banking77">banking77</a> dataset.
<br><br>
<b>What it does.</b> It <i>sorts</i> messages, it does not answer them. A real support
desk would use this to route each incoming message to the right team in ~10 ms rather
than having a person read it just to decide where it goes.
<br><br>
<b>What it isn't.</b> Not connected to any bank. It cannot see or act on any account.
"""

EXAMPLES = [
    "My card still hasn't arrived after two weeks, what should I do?",
    "I lost my card, someone might be using it",
    "The ATM ate my card and didn't give it back",
    "My money got taken out twice for the same purchase",
    "Why is there an extra charge on my statement?",
    "How long does a transfer to another country take?",
    "What's the exchange rate you use?",
    "I want to close my account",
]


def load_metrics(mode: str) -> dict:
    """Read a mode's training metrics so the UI can show real measured numbers.

    Returns an empty dict rather than failing — the demo must still run for
    someone who cloned the repo and has only a checkpoint.
    """
    path = RESULTS_DIR / f"{mode}_metrics.json"
    if not path.exists():
        return {}
    try:
        return json.loads(path.read_text())
    except (json.JSONDecodeError, OSError):
        return {}


class IntentPredictor:
    """One trained model plus its tokenizer, loaded once at startup.

    Loading inside the predict function instead would re-read the checkpoint on
    every request, turning a ~9 ms inference into a multi-second one. On a
    Hugging Face Space that is the most common cause of "why is my demo so slow".
    """

    def __init__(self, checkpoint: Path, device: str = "cpu", label: str = "") -> None:
        if not checkpoint.exists():
            raise FileNotFoundError(
                f"No checkpoint at {checkpoint}. Train one first:\n"
                f"    python train.py --mode lora"
            )
        self.device = torch.device(device)
        self.checkpoint = checkpoint
        self.model = TextClassifier.load(checkpoint, device=self.device)
        self.model.eval()
        self.tokenizer = AutoTokenizer.from_pretrained(self.model.model_name)

        # Class names come from the checkpoint, not from the dataset or a
        # hardcoded list. If they were re-derived here and the order differed by
        # even one position, every prediction would display the wrong name while
        # looking perfectly plausible.
        self.labels = self.model.label_names or [
            f"class_{i}" for i in range(self.model.num_labels)
        ]

        report = self.model.trainable_parameter_report()
        self.trainable = int(report["trainable_params"])
        self.trainable_pct = float(report["trainable_pct"])
        self.size_kb = checkpoint.stat().st_size / 1024
        self.display = label or str(report["mode"])
        self.train_seconds = load_metrics(str(report["mode"])).get("train_seconds")

        log.info("Loaded %s: %s trainable params (%.3f%%), %.0f KB",
                 self.display, f"{self.trainable:,}", self.trainable_pct, self.size_kb)

    def predict(self, text: str) -> tuple[dict[str, float], float]:
        """Classify one query.

        Returns:
            ``({"card arrival": 0.99, ...}, latency_ms)``. Probabilities are a
            softmax over all 77 classes and sum to 1.
        """
        if not text or not text.strip():
            # Uniform reads honestly as "no input", rather than whatever the
            # model happens to emit for an empty string.
            return {name: 1.0 / len(self.labels) for name in self.labels}, 0.0

        encoded = self.tokenizer(
            text.strip(), truncation=True, max_length=MAX_LENGTH,
            return_tensors="pt",  # (1, L) tensors rather than plain lists
        )
        start = time.perf_counter()
        _, probabilities = self.model.predict(
            encoded["input_ids"].to(self.device),
            encoded["attention_mask"].to(self.device),
        )
        latency_ms = 1000 * (time.perf_counter() - start)

        # (1, 77) -> a plain dict. card_arrival -> "card arrival" for display.
        return (
            {self.labels[i].replace("_", " "): float(p)
             for i, p in enumerate(probabilities[0])},
            latency_ms,
        )

    def size_str(self) -> str:
        """Checkpoint size, in whichever unit reads better."""
        return (f"{self.size_kb / 1024:.0f} MB" if self.size_kb > 1024
                else f"{self.size_kb:.0f} KB")

    def subtitle(self) -> str:
        """Compact cost summary shown under each model's panel heading."""
        parts = [f"{self.trainable:,} params", self.size_str()]
        if self.train_seconds:
            parts.append(f"{self.train_seconds:.0f}s to train")
        return " · ".join(parts)


class ComparisonPredictor:
    """Runs both models on the same input so they can be compared live."""

    def __init__(self, lora: IntentPredictor, full: IntentPredictor | None) -> None:
        self.lora = lora
        self.full = full

    def predict(self, text: str) -> tuple[dict[str, float], dict[str, float], str]:
        """Classify with both models and describe how they compare.

        Returns ``(lora_probs, full_probs, verdict_markdown)``.
        """
        lora_probs, lora_ms = self.lora.predict(text)
        if self.full is None:
            return lora_probs, {}, ""

        full_probs, full_ms = self.full.predict(text)
        if not text or not text.strip():
            return lora_probs, full_probs, "_Enter a question above._"

        lora_top = max(lora_probs, key=lora_probs.get)
        full_top = max(full_probs, key=full_probs.get)
        agree = lora_top == full_top

        # Per-query facts only. Anything constant across queries (parameter
        # counts, checkpoint size, training time) lives in the static cost table
        # below, so this does not repeat itself on every click.
        if agree:
            head = f"Both models predict <b>{lora_top}</b>"
            detail = (
                f"LoRA {lora_probs[lora_top]:.0%} confident in {lora_ms:.0f} ms · "
                f"full fine-tuning {full_probs[full_top]:.0%} in {full_ms:.0f} ms. "
                f"Same answer from a model that trained "
                f"<b>{self.full.trainable / self.lora.trainable:,.0f}× fewer parameters</b>."
            )
        else:
            head = (f"The models disagree — LoRA says <b>{lora_top}</b>, "
                    f"full fine-tuning says <b>{full_top}</b>")
            detail = (
                f"{lora_probs[lora_top]:.0%} vs {full_probs[full_top]:.0%} confidence. "
                "Disagreements are usually genuinely ambiguous queries; across the whole "
                "test set the two land within 0.4 percentage points of each other."
            )
        verdict = (f'<div class="verdict"><div class="head">{head}</div>'
                   f'<div class="detail">{detail}</div></div>')
        return lora_probs, full_probs, verdict


def load_test_accuracy() -> dict[str, float]:
    """Pull each mode's test accuracy out of ``results/comparison.csv``.

    Parsed rather than hardcoded so the UI cannot drift from the last real
    evaluation. Returns ``{}`` if the file is absent, and the caller omits the
    row rather than showing a wrong number.
    """
    path = RESULTS_DIR / "comparison.csv"
    if not path.exists():
        return {}
    try:
        lines = path.read_text().strip().splitlines()
        header = lines[0].split(",")
        mode_i, acc_i = header.index("model"), header.index("test_accuracy")
        out = {}
        for line in lines[1:]:
            cells = line.split(",")
            if cells[mode_i] in ("lora", "full") and cells[acc_i]:
                out[cells[mode_i]] = float(cells[acc_i])
        return out
    except (OSError, ValueError, IndexError):
        return {}


def cost_table(lora: IntentPredictor, full: IntentPredictor) -> str:
    """Static HTML table: what each model cost to train, store, and run.

    These numbers do not change per query, so they render once beneath the live
    comparison rather than being recomputed on every click. The LoRA column is
    highlighted wherever it wins, which is every row except accuracy — and it
    edges that one too.
    """
    rows = [
        ("Trainable parameters", f"{lora.trainable:,}", f"{full.trainable:,}",
         f"{full.trainable / lora.trainable:,.0f}× fewer", True),
        ("Share of the model", f"{lora.trainable_pct:.2f}%",
         f"{full.trainable_pct:.0f}%", "", True),
        ("Checkpoint size", lora.size_str(), full.size_str(),
         f"{full.size_kb / lora.size_kb:,.0f}× smaller", True),
    ]
    if lora.train_seconds and full.train_seconds:
        rows.append((
            "Training time", f"{lora.train_seconds:.0f}s", f"{full.train_seconds:.0f}s",
            f"{full.train_seconds / lora.train_seconds:.1f}× faster", True,
        ))

    accuracy = load_test_accuracy()
    if "lora" in accuracy and "full" in accuracy:
        delta = (accuracy["lora"] - accuracy["full"]) * 100
        rows.insert(0, ("Test accuracy", f"{accuracy['lora']:.1%}",
                        f"{accuracy['full']:.1%}", f"{delta:+.1f} pts",
                        accuracy["lora"] >= accuracy["full"]))

    body = "".join(
        f'<tr><td>{name}</td>'
        f'<td class="{"win" if win else ""}">{a}</td>'
        f'<td>{b}</td><td>{note}</td></tr>'
        for name, a, b, note, win in rows
    )
    return (
        '<div class="costwrap"><table class="cost">'
        "<thead><tr><th></th><th>LoRA (r=8)</th><th>Full fine-tuning</th>"
        "<th></th></tr></thead>"
        f"<tbody>{body}</tbody></table></div>"
        '<div class="foot" style="margin-top:.6rem">'
        "Accuracy is on the held-out test split. Inference latency is deliberately "
        "absent: LoRA does not improve it. Both models run the identical "
        "66M-parameter forward pass — LoRA saves training cost and storage, not "
        "inference time.</div>"
    )


#: All colours come from Gradio's own theme variables rather than fixed hex
#: values, so the page reads correctly in both light and dark mode — visitors can
#: toggle, and hardcoded colours would break one of the two.
CSS = """
.hero { text-align: center; padding: 1.5rem 0 0.5rem; }
.hero h1 { font-size: 2.6rem; margin: 0 0 .3rem; letter-spacing: -0.02em; }
.hero p { color: var(--body-text-color-subdued); margin: 0 auto; max-width: 42rem;
          font-size: 1.02rem; line-height: 1.55; }

.statrow { display: flex; gap: .75rem; justify-content: center; flex-wrap: wrap;
           margin: 1.25rem 0 .5rem; }
.stat { flex: 1 1 8.5rem; min-width: 8.5rem; max-width: 12rem; padding: .85rem .75rem;
        border: 1px solid var(--border-color-primary); border-radius: 10px;
        background: var(--background-fill-secondary); text-align: center; }
.stat .v { font-size: 1.5rem; font-weight: 700; line-height: 1.15;
           font-variant-numeric: tabular-nums; }
.stat .k { font-size: .72rem; text-transform: uppercase; letter-spacing: .06em;
           color: var(--body-text-color-subdued); margin-top: .3rem; }
.stat.accent .v { color: var(--primary-500); }

.panel { border: 1px solid var(--border-color-primary); border-radius: 12px;
         padding: .9rem 1rem 1rem; background: var(--background-fill-primary); }
.panel.win { border-color: var(--primary-500); }
.panel h3 { margin: 0 0 .15rem; font-size: 1.05rem; }
.panel .sub { font-size: .8rem; color: var(--body-text-color-subdued);
              font-variant-numeric: tabular-nums; }

.verdict { border-radius: 10px; padding: .75rem 1rem; margin-top: .25rem;
           border: 1px solid var(--border-color-primary);
           background: var(--background-fill-secondary); }
.verdict .head { font-weight: 650; margin-bottom: .2rem; }
.verdict .detail { font-size: .88rem; color: var(--body-text-color-subdued);
                   line-height: 1.5; }

.costwrap { overflow-x: auto; }
.cost { width: 100%; border-collapse: collapse; font-size: .9rem; }
.cost th, .cost td { padding: .5rem .7rem; text-align: right;
                     border-bottom: 1px solid var(--border-color-primary);
                     font-variant-numeric: tabular-nums; white-space: nowrap; }
.cost th:first-child, .cost td:first-child { text-align: left; white-space: normal; }
.cost thead th { color: var(--body-text-color-subdued); font-weight: 600;
                 text-transform: uppercase; font-size: .72rem; letter-spacing: .05em; }
.cost td.win { color: var(--primary-500); font-weight: 650; }
.foot { color: var(--body-text-color-subdued); font-size: .85rem; line-height: 1.6; }
.foot a { color: var(--primary-500); }
"""


def hero_html(lora: IntentPredictor, full: IntentPredictor | None) -> str:
    """Headline block: what this is, plus the four numbers that are the point."""
    accuracy = load_test_accuracy()
    cards = []
    if "lora" in accuracy:
        cards.append(("accent", f"{accuracy['lora']:.1%}", "Test accuracy"))
    cards.append(("accent", f"{lora.trainable_pct:.2f}%", "Of params trained"))
    if full:
        cards.append(("", f"{full.trainable / lora.trainable:,.0f}×", "Fewer params"))
        cards.append(("", f"{full.size_kb / lora.size_kb:,.0f}×", "Smaller file"))
    else:
        cards.append(("", f"{lora.trainable:,}", "Trainable params"))
        cards.append(("", f"{lora.size_kb:,.0f} KB", "Checkpoint"))

    stats = "".join(
        f'<div class="stat {cls}"><div class="v">{v}</div><div class="k">{k}</div></div>'
        for cls, v, k in cards
    )
    tagline = (
        "Two DistilBERT models classify the same customer banking query into one of "
        f"<b>{len(lora.labels)} support intents</b> — one fine-tuned normally, one with "
        "<b>LoRA implemented from scratch</b>. Same answers, a fraction of the training."
        if full else
        f"DistilBERT routes a customer banking query to one of <b>{len(lora.labels)} "
        "support intents</b>, fine-tuned with <b>LoRA implemented from scratch</b>."
    )
    return (
        f'<div class="hero"><h1>RiscAutious</h1><p>{tagline}</p></div>'
        f'<div class="statrow">{stats}</div>'
    )


def build_interface(predictor: ComparisonPredictor) -> gr.Blocks:
    """Assemble the Gradio UI."""
    lora, full = predictor.lora, predictor.full

    # Gradio 6 moved `theme` and `css` from Blocks() to launch(); they are
    # applied in main() rather than here.
    with gr.Blocks(title="RiscAutious — LoRA vs full fine-tuning") as demo:
        gr.HTML(hero_html(lora, full))

        with gr.Group():
            text_input = gr.Textbox(
                label="Ask what a bank customer would ask",
                placeholder="My card hasn't arrived yet…",
                lines=2,
            )
            submit = gr.Button("Classify", variant="primary", size="lg")

        gr.Examples(examples=EXAMPLES, inputs=text_input, label="Or try one of these")

        # Outcomes first, side by side — the two predictions are what a visitor
        # came to see. The cost tables that justify them go underneath.
        full_out = None
        with gr.Row(equal_height=True):
            with gr.Column():
                with gr.Column(elem_classes="panel win"):
                    gr.HTML(f'<h3>LoRA <span class="sub">r=8</span></h3>'
                            f'<div class="sub">{lora.subtitle()}</div>')
                    # 77 classes is far too many to show; display the top few.
                    lora_out = gr.Label(label="Predicted intent",
                                        num_top_classes=TOP_K, show_label=False)
            if full:
                with gr.Column():
                    with gr.Column(elem_classes="panel"):
                        gr.HTML(f'<h3>Full fine-tuning</h3>'
                                f'<div class="sub">{full.subtitle()}</div>')
                        full_out = gr.Label(label="Predicted intent",
                                            num_top_classes=TOP_K, show_label=False)

        # Metrics below the outcomes: the live per-query comparison first, then
        # the static cost that does not change between queries.
        verdict = gr.HTML()
        if full:
            gr.HTML(cost_table(lora, full))

        gr.HTML(f'<div class="foot">{DISCLAIMER}</div>')

        if full:
            outputs: list = [lora_out, full_out, verdict]
            handler = predictor.predict
        else:
            outputs = [lora_out]
            handler = lambda text: predictor.lora.predict(text)[0]  # noqa: E731

        # Fire on button click and on Enter, so the demo feels responsive.
        submit.click(handler, inputs=text_input, outputs=outputs)
        text_input.submit(handler, inputs=text_input, outputs=outputs)

    return demo


def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
    """Define and parse the command-line interface."""
    p = argparse.ArgumentParser(description="Serve the banking intent classifier demo.")
    p.add_argument("--checkpoint", type=Path, default=DEFAULT_LORA,
                   help="LoRA checkpoint (the one the demo is built around).")
    p.add_argument("--full-checkpoint", type=Path, default=DEFAULT_FULL,
                   help="Full fine-tuned model, for the side-by-side comparison.")
    p.add_argument("--lora-only", action="store_true",
                   help="Skip the 253 MB full model — what a Space should do.")
    p.add_argument("--device", default="cpu",
                   help="CPU is right for a demo; the model is small and requests single.")
    p.add_argument("--share", action="store_true", help="Create a public gradio.live link.")
    p.add_argument("--port", type=int, default=7860)
    return p.parse_args(argv)


def main(argv: Sequence[str] | None = None) -> int:
    """Entry point. Returns a process exit code."""
    logging.basicConfig(level=logging.INFO, format="%(levelname)-7s %(message)s")
    for noisy in ("httpx", "urllib3", "filelock", "huggingface_hub"):
        logging.getLogger(noisy).setLevel(logging.WARNING)
    import transformers
    transformers.logging.set_verbosity_error()

    args = parse_args(argv)
    try:
        lora = IntentPredictor(args.checkpoint, args.device, "LoRA (r=8)")
    except FileNotFoundError as exc:
        log.error("%s", exc)
        return 1

    full = None
    if args.lora_only:
        log.info("--lora-only: skipping the full fine-tuned model")
    elif args.full_checkpoint.exists():
        full = IntentPredictor(args.full_checkpoint, args.device, "Full fine-tuning")
    else:
        log.info("No %s — showing LoRA only. Train it with: "
                 "python train.py --mode full", args.full_checkpoint)

    # Keep launch() kwargs minimal — Gradio changes these between major versions
    # (6 removed `show_api`), and a Space failing to start over a cosmetic flag
    # is a bad trade.
    theme = gr.themes.Soft(
        primary_hue="indigo", neutral_hue="slate",
        font=[gr.themes.GoogleFont("Inter"), "system-ui", "sans-serif"],
    )
    build_interface(ComparisonPredictor(lora, full)).launch(
        server_port=args.port, share=args.share, theme=theme, css=CSS
    )
    return 0


if __name__ == "__main__":
    sys.exit(main())