app.py

"""
Smart MCQ Solver — Gradio Web Application
Loads ELECTRA-base + LoRA fine-tuned model weights from Hugging Face Hub (SpreadSheets600/electra-lora-mcq)
Scored by MAP@3 on Kaggle Science MCQ (0.8849 OOF MAP@3)
"""

import time
import threading
import numpy as np
import torch
import torch.nn as nn
import gradio as gr
from transformers import AutoTokenizer, AutoModel
from peft import PeftModel

try:
    import spaces
    HAS_SPACES = True
except ImportError:
    HAS_SPACES = False

# ── Model Configuration & Hub Repositories ────────────────────────────────
ADAPTER_REPO = "SpreadSheets600/electra-lora-mcq"
BACKBONE     = "google/electra-base-discriminator"
LABELS       = ["A", "B", "C", "D", "E"]
MAX_LENGTH   = 256

# ── Neural Architecture (mirrors notebook) ─────────────────────────────────
class EncoderMCQModel(nn.Module):
    def __init__(self, backbone):
        super().__init__()
        self.encoder    = backbone
        h               = backbone.config.hidden_size
        self.dropouts   = nn.ModuleList([nn.Dropout(0.2) for _ in range(5)])
        self.classifier = nn.Linear(h, 1)

    def mean_pool(self, hidden, mask):
        m = mask.unsqueeze(-1).float()
        return (hidden * m).sum(1) / m.sum(1).clamp(min=1e-9)

    def forward(self, input_ids, attention_mask):
        b, n, L = input_ids.shape
        out  = self.encoder(
            input_ids=input_ids.view(b*n, L),
            attention_mask=attention_mask.view(b*n, L),
        )
        pool  = self.mean_pool(out.last_hidden_state, attention_mask.view(b*n, L))
        logit = torch.mean(
            torch.stack([self.classifier(dp(pool)) for dp in self.dropouts]), 0
        )
        return logit.view(b, n)

# ── Lazy Thread-Safe Model Loader ──────────────────────────────────────────
_STATE = {"tokenizer": None, "model": None, "ready": False, "error": None}
_LOCK  = threading.Lock()

def _load_model():
    with _LOCK:
        if _STATE["ready"] or _STATE["error"]:
            return
        try:
            tok  = AutoTokenizer.from_pretrained(BACKBONE)
            base = AutoModel.from_pretrained(BACKBONE)
            wrap = EncoderMCQModel(base)
            model = PeftModel.from_pretrained(wrap, ADAPTER_REPO, is_trainable=False)
            model = model.eval()
            _STATE["tokenizer"] = tok
            _STATE["model"]     = model
            _STATE["ready"]     = True
        except Exception as e:
            _STATE["error"] = str(e)

def _ensure_model_ready():
    if not _STATE["ready"] and not _STATE["error"]:
        _load_model()
    if _STATE["error"]:
        raise RuntimeError(f"Model weight load error: {_STATE['error']}")

# ── Neural Inference Routine ───────────────────────────────────────────────
def run_electra_core(prompt, options):
    device = "cuda" if torch.cuda.is_available() else "cpu"
    tok   = _STATE["tokenizer"]
    model = _STATE["model"].to(device)

    ids_list, mask_list = [], []
    for opt in options:
        enc = tok(
            f"{prompt} [OPTION] {opt}",
            max_length=MAX_LENGTH,
            padding="max_length",
            truncation=True,
            return_tensors="pt"
        )
        ids_list.append(enc["input_ids"].squeeze(0))
        mask_list.append(enc["attention_mask"].squeeze(0))

    ids   = torch.stack(ids_list).unsqueeze(0).to(device)
    masks = torch.stack(mask_list).unsqueeze(0).to(device)

    logits = model(ids, masks).squeeze(0).cpu().float().numpy()
    probs  = torch.softmax(torch.tensor(logits), dim=0).numpy()
    return logits, probs

if HAS_SPACES:
    @spaces.GPU
    @torch.inference_mode()
    def run_electra_inference(prompt, options):
        return run_electra_core(prompt, options)
else:
    @torch.inference_mode()
    def run_electra_inference(prompt, options):
        return run_electra_core(prompt, options)

def calculate_top3(logits):
    return [LABELS[i] for i in np.argsort(-logits)[:3]]

def calculate_map3(correct_label, ranked_labels):
    for k, label in enumerate(ranked_labels[:3], 1):
        if label == correct_label:
            return 1.0 / k
    return 0.0

# ── Gradio Prediction Callback ─────────────────────────────────────────────
def predict_mcq(prompt, opt_a, opt_b, opt_c, opt_d, opt_e, correct_key):
    if not prompt or not prompt.strip():
        raise gr.Error("Please enter a valid question prompt.")

    options = [opt_a, opt_b, opt_c, opt_d, opt_e]
    if any(not o or not o.strip() for o in options):
        raise gr.Error("Please fill in all 5 option choices (A–E).")

    _ensure_model_ready()

    t0 = time.perf_counter()
    logits, probs = run_electra_inference(prompt.strip(), [o.strip() for o in options])
    elapsed_ms = (time.perf_counter() - t0) * 1000

    ranked = calculate_top3(logits)
    medals = ["🥇 1st", "🥈 2nd", "🥉 3rd"]

    top3_markdown = "   |   ".join([f"{medals[r]}: **{ranked[r]}**" for r in range(3)])

    conf_dict = {f"Option {LABELS[i]}": float(probs[i]) for i in range(5)}

    map3_markdown = ""
    if correct_key and correct_key.strip().upper() in LABELS:
        cl = correct_key.strip().upper()
        map3_val = calculate_map3(cl, ranked)
        pos = next((r + 1 for r, l in enumerate(ranked) if l == cl), None)
        pos_str = f"rank **{pos}**" if pos else "outside top-3"
        map3_markdown = f"📊 MAP@3 Metric = **{map3_val:.4f}** (Correct Answer: **{cl}** found at {pos_str})"

    order = np.argsort(-logits)
    table_rows = []
    for rank_idx, idx in enumerate(order):
        lbl = LABELS[idx]
        tick = " ✓" if correct_key and lbl == correct_key.strip().upper() else ""
        table_rows.append([
            f"#{rank_idx + 1}",
            f"{lbl}{tick}",
            options[idx][:70] + ("…" if len(options[idx]) > 70 else ""),
            f"{probs[idx] * 100:.1f}%",
            f"{logits[idx]:.3f}"
        ])

    device_used = "ZeroGPU" if torch.cuda.is_available() else "CPU"
    info_markdown = (
        f"⏱ **{elapsed_ms:.0f} ms** response time  ·  "
        f"Device: **{device_used}**  ·  "
        f"Model: **ELECTRA-base + LoRA ($r=16$)**  ·  "
        f"5-Fold OOF MAP@3: **0.8849**"
    )

    return top3_markdown, conf_dict, table_rows, map3_markdown, info_markdown

# ── Pre-loaded Test Cases ──────────────────────────────────────────────────
EXAMPLES = [
    ["Which of the following is a prime number?", "15", "21", "23", "25", "27", "C"],
    ["What is the chemical symbol for gold?", "Ag", "Au", "Fe", "Pb", "Hg", "B"],
    ["The process by which plants make food using sunlight is called:", "Respiration", "Fermentation", "Photosynthesis", "Transpiration", "Osmosis", "C"],
    ["Which planet in our solar system has the most moons?", "Jupiter", "Saturn", "Uranus", "Neptune", "Mars", "B"],
    ["DNA replication occurs during which phase of the cell cycle?", "G1 phase", "S phase", "G2 phase", "M phase", "G0 phase", "B"],
    ["What is the powerhouse of the cell?", "Nucleus", "Ribosome", "Mitochondria", "Golgi apparatus", "Endoplasmic reticulum", "C"]
]

# ── Gradio Blocks Interface ────────────────────────────────────────────────
with gr.Blocks(title="Smart MCQ Solver") as demo:

    gr.HTML("""
    <div style="text-align:center; padding:20px 0 10px">
      <h1 style="margin:0; font-size:2.3rem; font-weight:800; color:#4f46e5">
        🎯 Smart MCQ Solver
      </h1>
      <p style="color:#475569; font-size:1.05rem; margin:6px 0 0">
        ELECTRA-base + LoRA Neural Option Scorer &nbsp;·&nbsp;
        <strong style="color:#059669">MAP@3 = 0.8849</strong> on 5-Fold CV
      </p>
    </div>
    """)

    with gr.Row():
        with gr.Column(scale=5):
            gr.Markdown("### 📝 Question Input")
            prompt_input = gr.Textbox(
                label="Question Prompt",
                placeholder="Type your prompt here...",
                lines=3
            )
            with gr.Row():
                opt_a = gr.Textbox(label="Option A")
                opt_b = gr.Textbox(label="Option B")
            with gr.Row():
                opt_c = gr.Textbox(label="Option C")
                opt_d = gr.Textbox(label="Option D")
            opt_e = gr.Textbox(label="Option E")
            correct_input = gr.Textbox(
                label="✅ Correct Answer Key (Optional - A/B/C/D/E for MAP@3 Evaluation)",
                placeholder="Leave blank if unknown",
                max_lines=1
            )
            with gr.Row():
                predict_button = gr.Button("🔍 Predict Top-3 & Confidence Scores", variant="primary", size="lg")
                clear_button = gr.ClearButton(
                    [prompt_input, opt_a, opt_b, opt_c, opt_d, opt_e, correct_input],
                    value="🗑 Clear",
                    size="lg"
                )

        with gr.Column(scale=5):
            gr.Markdown("### 🏆 Model Predictions & Analytics")
            top3_output = gr.Markdown(value="*Submit a question to see predictions.*")
            map3_output = gr.Markdown()

            gr.Markdown("#### Option Confidence Distribution")
            conf_output = gr.Label(label="Softmax Probability", num_top_classes=5)

            gr.Markdown("#### Full Option Ranking Details")
            table_output = gr.Dataframe(
                headers=["Rank", "Option", "Text", "Confidence %", "Logit"],
                interactive=False,
                wrap=True
            )
            info_output = gr.Markdown()

    gr.Markdown("---")
    gr.Examples(
        examples=EXAMPLES,
        inputs=[prompt_input, opt_a, opt_b, opt_c, opt_d, opt_e, correct_input],
        outputs=[top3_output, conf_output, table_output, map3_output, info_output],
        fn=predict_mcq,
        cache_examples=False,
        label="📚 Load Example Test Case",
        examples_per_page=6
    )

    with gr.Accordion("ℹ️ Model Architecture & Evaluation Specs", open=False):
        gr.Markdown("""
| Component | Specification |
|---|---|
| **Backbone Model** | `google/electra-base-discriminator` |
| **LoRA Adapter** | `SpreadSheets600/electra-lora-mcq` ($r=16, \\alpha=32$) |
| **Trainable Parameters** | ~1.18M (~1% of total weights) |
| **Scoring Head** | Joint 5-way scoring via 5× Dropout(0.2) + Linear layer |
| **5-Fold CV MAP@3** | **0.8849** |
| **5-Fold CV Accuracy** | **81.1%** |
| **Course & Student** | BITS Pilani WILP (DL & GenAI) · Student 24F2008474 |
        """)

    predict_button.click(
        fn=predict_mcq,
        inputs=[prompt_input, opt_a, opt_b, opt_c, opt_d, opt_e, correct_input],
        outputs=[top3_output, conf_output, table_output, map3_output, info_output]
    )

try:
    demo.launch(ssr_mode=False)
except TypeError:
    demo.launch()