#!/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 = """
What this is. A demonstration of LoRA implemented from scratch — no
peft — on the public
banking77 dataset.
What it does. It sorts 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.
What it isn't. 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 {lora_top}"
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"{self.full.trainable / self.lora.trainable:,.0f}× fewer parameters."
)
else:
head = (f"The models disagree — LoRA says {lora_top}, "
f"full fine-tuning says {full_top}")
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'')
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'| {name} | '
f'{a} | '
f'{b} | {note} |
'
for name, a, b, note, win in rows
)
return (
''
" | LoRA (r=8) | Full fine-tuning | "
" |
|---|
"
f"{body}
{stats}
'
)
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'LoRA r=8
'
f'{lora.subtitle()}
')
# 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'Full fine-tuning
'
f'{full.subtitle()}
')
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'')
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())