Spaces:
Sleeping
Sleeping
| #!/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()) | |