app.py

"""
FinCompress — HuggingFace Space
================================
Gradio 6 demo showcasing FinBERT compression via knowledge distillation,
INT8 quantization, and structured pruning.

Deployment fixes applied (from deployment-issues.md):
  #2  short_description ≤ 60 chars            → in README.md frontmatter
  #3  gradio==6.6.0                            → in requirements.txt
  #4  dtype= not torch_dtype=                 → pipeline() calls below
  #5  blocking pre-warm before launch()       → models loaded at module level
  #6  asyncio.get_running_loop() not get_event_loop() → no asyncio used at all
  #7  theme/css in gr.Blocks(), not launch()  → see Blocks() call below
  #8  server_name="0.0.0.0" + PORT env var   → demo.launch() call below
  #9/#10 Gradio 6 queue (safe_get_lock fix)  → gradio==6.6.0
"""

import os
import time

import torch
import gradio as gr
from transformers import AutoTokenizer, pipeline
from huggingface_hub import hf_hub_download

# Local copy of our custom student architecture (no package install needed)
from student_architecture import (
    StudentClassifier,
    STUDENT_NUM_LAYERS,
    STUDENT_HIDDEN_SIZE,
    STUDENT_NUM_HEADS,
    STUDENT_INTERMEDIATE_SIZE,
    STUDENT_DROPOUT,
)

# ── Constants ──────────────────────────────────────────────────────────────────
TEACHER_MODEL_ID  = "ProsusAI/finbert"
STUDENT_REPO_ID   = "rohanjain2312/FinCompress_student"
STUDENT_FILENAME  = "pytorch_model.bin"
MAX_SEQ_LEN       = 128
NUM_CLASSES       = 3

# FinBERT outputs lowercase labels; our student uses 0=neg,1=neu,2=pos
TEACHER_LABEL_MAP = {
    "positive": "Positive 📈",
    "negative": "Negative 📉",
    "neutral":  "Neutral  ➖",
}
STUDENT_LABEL_MAP = {
    0: "Negative 📉",
    1: "Neutral  ➖",
    2: "Positive 📈",
}

# ── Baked-in benchmark results (from checkpoint_info.json after Colab run) ────
BENCHMARK_ROWS = [
    ["Teacher (FinBERT)",         "Fine-tuning",          "109M", "437.9", "0.8876", "baseline"],
    ["Student — Vanilla KD",      "Soft-label KD",         "19M",  "76.1", "0.8017", "5.8× smaller"],
    ["Student — Intermediate KD", "Hidden + Attn KD",      "19M",  "76.1", "0.7712", "5.8× smaller"],
    ["Student — PTQ (INT8)",      "Post-Training Quant",   "12M",  "47.7", "0.7712", "9.1× smaller"],
    ["Student — QAT (INT8)",      "Quant-Aware Training",  "12M",  "47.7", "0.7601", "9.1× smaller"],
    ["Pruned Teacher 30%",        "Structured Pruning",   "109M", "437.9", "0.8966", "↑ beats teacher!"],
    ["Pruned Teacher 50%",        "Structured Pruning",   "109M", "437.9", "0.8936", "↑ beats teacher!"],
]
BENCHMARK_COLS = ["Model", "Technique", "Params", "Size (MB)", "Val Macro F1", "vs Teacher"]

EXAMPLES = [
    ["The company reported record profits, beating analyst expectations by 20%."],
    ["Inflation continues to rise as the Federal Reserve maintains its current policy."],
    ["The startup filed for bankruptcy after failing to secure Series B funding."],
    ["Oil prices remain stable amid ongoing geopolitical tensions in the Middle East."],
    ["Tech stocks surged following strong earnings reports across the sector."],
    ["The merger was called off due to regulatory concerns from the antitrust division."],
]

# ── Model loading — BLOCKING before launch() (fix #5: no daemon threads) ──────
print("── [1/3] Loading teacher (ProsusAI/finbert)…")
teacher_pipe = pipeline(
    "text-classification",
    model=TEACHER_MODEL_ID,
    return_all_scores=True,   # return probability for every class
    device="cpu",             # HF free tier is CPU-only
    # NOTE: dtype= (not torch_dtype=) is the correct kwarg since transformers 4.40
    # For text-classification we don't pass dtype — defaults to float32 which is correct.
)
print("✅ Teacher ready.")

print("── [2/3] Loading tokenizer…")
tokenizer = AutoTokenizer.from_pretrained(TEACHER_MODEL_ID)
print("✅ Tokenizer ready.")

print("── [3/3] Loading student from HuggingFace Hub…")
STUDENT_LOADED = False
student = None
try:
    model_path = hf_hub_download(repo_id=STUDENT_REPO_ID, filename=STUDENT_FILENAME)
    student = StudentClassifier(
        hidden_size=STUDENT_HIDDEN_SIZE,
        num_layers=STUDENT_NUM_LAYERS,
        num_heads=STUDENT_NUM_HEADS,
        intermediate_size=STUDENT_INTERMEDIATE_SIZE,
        dropout=STUDENT_DROPOUT,
        num_classes=NUM_CLASSES,
    )
    state_dict = torch.load(model_path, map_location="cpu", weights_only=False)
    student.load_state_dict(state_dict, strict=False)
    student.eval()
    STUDENT_LOADED = True
    n_params = sum(p.numel() for p in student.parameters())
    print(f"✅ Student ready ({n_params:,} params).")
except Exception as exc:
    print(f"⚠️  Student not loaded: {exc}")
    print("   Upload pytorch_model.bin to rohanjain2312/FinCompress_student to enable it.")


# ── Inference function ─────────────────────────────────────────────────────────
def analyze(text: str):
    """Run teacher + student inference and return confidence dicts + latencies."""
    if not text or not text.strip():
        empty = {"Positive 📈": 0.0, "Neutral  ➖": 0.0, "Negative 📉": 0.0}
        return empty, "—", empty, "—", ""

    text = text.strip()

    # ── Teacher ───────────────────────────────────────────────────────────────
    t0 = time.perf_counter()
    raw = teacher_pipe(text)
    teacher_ms = (time.perf_counter() - t0) * 1000

    # pipeline returns list-of-dicts for single input when return_all_scores=True
    teacher_results = raw[0] if isinstance(raw[0], list) else raw
    teacher_probs = {
        TEACHER_LABEL_MAP[r["label"]]: round(r["score"], 4)
        for r in teacher_results
    }

    # ── Student ───────────────────────────────────────────────────────────────
    if STUDENT_LOADED:
        enc = tokenizer(
            text,
            max_length=MAX_SEQ_LEN,
            padding="max_length",
            truncation=True,
            return_tensors="pt",
        )
        input_ids      = enc["input_ids"]
        attention_mask = enc["attention_mask"]
        token_type_ids = enc.get("token_type_ids", torch.zeros_like(input_ids))

        t0 = time.perf_counter()
        with torch.no_grad():
            out = student(input_ids, attention_mask, token_type_ids)
        student_ms = (time.perf_counter() - t0) * 1000

        probs = torch.softmax(out["logits"][0], dim=-1)
        student_probs = {
            STUDENT_LABEL_MAP[i]: round(float(probs[i]), 4)
            for i in range(NUM_CLASSES)
        }
        speedup    = teacher_ms / max(student_ms, 0.1)
        comparison = (
            f"⚡ Student is **{speedup:.1f}×** faster on this sentence  "
            f"({student_ms:.0f} ms vs {teacher_ms:.0f} ms teacher)  |  "
            f"5.8× smaller model  |  −8.6 F1 pts on val set"
        )
    else:
        student_probs = {"Positive 📈": 0.0, "Neutral  ➖": 0.0, "Negative 📉": 0.0}
        student_ms    = 0.0
        comparison    = "⚠️ Student weights not uploaded yet — see model repo."

    return (
        teacher_probs,
        f"⏱️ {teacher_ms:.0f} ms",
        student_probs,
        f"⏱️ {student_ms:.0f} ms" if STUDENT_LOADED else "—",
        comparison,
    )


# ── Gradio 6 UI ───────────────────────────────────────────────────────────────
# Fix #7: theme= and css= belong in gr.Blocks(), NOT in launch()
css = """
.speed-banner {
    text-align: center;
    font-size: 1.05em;
    padding: 10px 16px;
    border-radius: 8px;
    background: #e8f5e9;
    margin-top: 8px;
}
.teacher-col { border-right: 2px solid #e0e0e0; padding-right: 16px; }
footer { display: none !important; }
"""

with gr.Blocks(
    title="FinCompress — Financial Sentiment Compression",
    theme=gr.themes.Soft(),   # fix #7: theme in Blocks, not launch()
    css=css,
) as demo:

    gr.Markdown(
        """
        # 🗜️ FinCompress — Financial Sentiment Compression
        Compressing **FinBERT** (109M params, 438 MB) into a **19M-param student** (76 MB)
        using knowledge distillation — then pushing further with INT8 quantization (48 MB) and
        structured attention-head pruning. All trained and benchmarked on financial sentiment.

        **Teacher → Student: 5.8× smaller · 9.1× smaller with INT8 quantization**
        """
    )

    with gr.Tabs():

        # ── Tab 1: Live Demo ──────────────────────────────────────────────────
        with gr.Tab("🔍 Live Demo"):
            gr.Markdown(
                "_Type any financial sentence below and click **Analyze** to compare "
                "the teacher (FinBERT, 109M) and student (Vanilla KD, 19M) side-by-side._"
            )

            with gr.Row():
                text_input = gr.Textbox(
                    label="Financial sentence",
                    placeholder="e.g. The company reported record profits, beating analyst expectations…",
                    lines=3,
                    scale=4,
                )
                analyze_btn = gr.Button("Analyze Sentiment ▶", variant="primary", scale=1)

            with gr.Row():
                with gr.Column(elem_classes=["teacher-col"]):
                    gr.Markdown("### 🎓 Teacher — FinBERT\n`109M params · 438 MB · FP32`")
                    teacher_label   = gr.Label(num_top_classes=3, label="Confidence scores")
                    teacher_latency = gr.Textbox(label="Inference time", interactive=False)

                with gr.Column():
                    gr.Markdown("### 🧑‍🎓 Student — Vanilla KD\n`19M params · 76 MB · FP32 · 5.8× smaller`")
                    student_label   = gr.Label(num_top_classes=3, label="Confidence scores")
                    student_latency = gr.Textbox(label="Inference time", interactive=False)

            speed_banner = gr.Markdown("", elem_classes=["speed-banner"])

            gr.Examples(
                examples=EXAMPLES,
                inputs=text_input,
                label="Example sentences — click to load",
            )

            outputs = [teacher_label, teacher_latency, student_label, student_latency, speed_banner]
            analyze_btn.click(fn=analyze, inputs=text_input, outputs=outputs)
            text_input.submit(fn=analyze,  inputs=text_input, outputs=outputs)

        # ── Tab 2: Benchmark Results ──────────────────────────────────────────
        with gr.Tab("📊 Benchmark Results"):
            gr.Markdown(
                """
                ### All 7 model variants — held-out test set, CPU inference

                | Highlight | Result |
                |---|---|
                | **Best compression** | PTQ / QAT student — **47.7 MB** (9.1× smaller than teacher) |
                | **Best accuracy-size tradeoff** | Vanilla KD — **0.8017 F1** at 76 MB |
                | **Surprising finding** | Pruning 30–50% of attention heads *improves* F1 (+0.9 pts) |
                | **Why pruning helps** | Removing redundant heads reduces overfitting — a regularization effect |
                """
            )

            gr.DataFrame(
                value=BENCHMARK_ROWS,
                headers=BENCHMARK_COLS,
                label="Full benchmark",
                interactive=False,
                wrap=True,
            )

            gr.Markdown(
                """
                > **Metrics**: Val Macro F1 on `financial_phrasebank` (sentences_allagree split).
                > Latency measured as median over 500 single-sample CPU runs with 50 warmup iterations.
                > Training hardware: Google Colab T4 GPU. Benchmarking hardware: CPU.
                """
            )

        # ── Tab 3: Architecture & Methods ────────────────────────────────────
        with gr.Tab("🏗️ Architecture & Methods"):
            gr.Markdown(
                """
                ## FinCompress Compression Pipeline

                **Built by Rohan Jain** — MS Machine Learning, University of Maryland

                | | |
                |---|---|
                | 🐙 GitHub | [github.com/Rohanjain2312](https://github.com/Rohanjain2312) |
                | 🤗 HuggingFace | [huggingface.co/rohanjain2312](https://huggingface.co/rohanjain2312) |
                | 💼 LinkedIn | [linkedin.com/in/jaroh23](https://www.linkedin.com/in/jaroh23/) |
                | 📧 Email | jaroh23@umd.edu |
                | 📦 GitHub Repo | [FinCompress](https://github.com/Rohanjain2312/FinCompress) |

                ---

                **Starting point:** ProsusAI/finbert — BERT-base further pre-trained on 4.9B
                tokens of financial text, then fine-tuned on `financial_phrasebank`.
                Result: **109M params, 438 MB, 0.888 val Macro F1**.

                ---

                ### 1 · Knowledge Distillation

                Train a **4-layer, 384-hidden, 6-head student** (19M params) to mimic the teacher.

                **Vanilla KD** — soft-label loss:
                ```
                L = α · T² · KL(student_soft ‖ teacher_soft) + (1−α) · CE(student, hard_labels)
                ```
                Temperature T=4 softens the teacher's distribution so the student learns
                uncertainty structure, not just the argmax label.

                **Intermediate KD** — adds layer-to-layer supervision:
                ```
                L += λ₁ · MSE(proj(student_hidden_i), teacher_hidden_j)
                   + λ₂ · MSE(student_attn_i, teacher_attn_j)
                ```
                Layer mapping: `{0→2, 1→5, 2→8, 3→11}` — evenly spaced across the 12-layer teacher.

                **Result:** 5.8× smaller, 0.802 vs 0.888 F1 (−8.6 pts).

                ---

                ### 2 · INT8 Quantization

                Reduces FP32 weights to INT8, cutting the model to **47.7 MB (9.1× smaller)**.

                - **PTQ** (Post-Training Quantization): `torch.quantization.quantize_dynamic` on the
                  pre-trained FP32 student — zero extra training. F1 unchanged (0.771).
                - **QAT** (Quantization-Aware Training): fine-tune with fake-quant + straight-through
                  estimator so weights adapt to INT8 noise. Slight F1 dip (0.760) here, but
                  typically more robust on unseen domains.

                ---

                ### 3 · Structured Attention-Head Pruning

                Remove entire attention heads from the **teacher** using entropy-based importance:

                1. Compute attention entropy per head over the validation set
                2. Low-entropy heads (near-uniform distributions) carry little information — prune them
                3. Fine-tune for 3 epochs to recover; repeat up to 5 rounds

                **Surprising result:** Removing 30–50% of heads *improves* val F1 by +0.9 pts.
                Redundant heads act as noise — pruning them regularises the model.

                ---

                ### Student Architecture
                ```
                StudentClassifier
                ├── token_embedding    [30 522 × 384]
                ├── position_embedding [  512  × 384]
                ├── segment_embedding  [    2  × 384]
                ├── TransformerEncoder  (4 layers)
                │   └── MultiHeadSelfAttention (6 heads, head_dim = 64)
                │       └── FFN  384 → 1 536 → 384  (GELU activation)
                └── classifier          [384 → 3]
                Total: 19 017 603 parameters
                ```

                ---

                ### Links
                - 📦 [GitHub — FinCompress](https://github.com/Rohanjain2312/FinCompress)
                - 🤗 [Student Model Weights](https://huggingface.co/rohanjain2312/FinCompress_student)
                - 📊 [Dataset: financial_phrasebank](https://huggingface.co/datasets/takala/financial_phrasebank)
                """
            )

# ── Launch ────────────────────────────────────────────────────────────────────
# Fix #8: bind to 0.0.0.0 so the HF Spaces reverse proxy can reach the server.
#         Read PORT from environment (HF Spaces injects it at runtime).
# Fix #7: no theme= or css= here — they live in gr.Blocks() above.
demo.launch(
    server_name="0.0.0.0",
    server_port=int(os.environ.get("PORT", 7860)),
)