""" Smart MCQ Solver — Gradio Web Application Loads 5-Fold ELECTRA-base fine-tuned model checkpoints from the local ./results directory. Each fold is a full state_dict of EncoderMCQModel (encoder backbone + scoring head). Predictions are ensembled across all available folds (average of logits). """ import os import glob import time import threading import numpy as np import torch import torch.nn as nn import gradio as gr from transformers import AutoTokenizer, AutoModel, AutoConfig try: import spaces HAS_SPACES = True except ImportError: HAS_SPACES = False # ── Model Configuration ------------------------------------------------ BACKBONE = "google/electra-base-discriminator" LABELS = ["A", "B", "C", "D", "E"] MAX_LENGTH = 256 # Directory containing the fold checkpoints ("results/electra_fold_*.pt") RESULTS_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "results") FOLD_FILES = sorted(glob.glob(os.path.join(RESULTS_DIR, "electra_fold_*.pt"))) # ── 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["models"] holds one EncoderMCQModel per available fold checkpoint. _STATE = {"tokenizer": None, "models": [], "ready": False, "error": None} _LOCK = threading.Lock() def _build_fold_model(state_dict): """Build a fresh EncoderMCQModel and load one fold's full state_dict into it.""" # The fold checkpoint contains the *entire* fine-tuned encoder, so we only # need the correct architecture (no need to download the pretrained weights). config = AutoConfig.from_pretrained(BACKBONE) base = AutoModel.from_config(config) model = EncoderMCQModel(base) model.load_state_dict(state_dict) # strict: overwrites encoder + classifier model.eval() return model def _load_model(): with _LOCK: if _STATE["ready"] or _STATE["error"]: return try: if not FOLD_FILES: raise FileNotFoundError( f"No `electra_fold_*.pt` checkpoints found in '{RESULTS_DIR}'. " "Please place your trained fold weights there." ) tok = AutoTokenizer.from_pretrained(BACKBONE) models = [] for fp in FOLD_FILES: ckpt = torch.load(fp, map_location="cpu", weights_only=False) models.append(_build_fold_model(ckpt)) del ckpt _STATE["tokenizer"] = tok _STATE["models"] = models _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"] models = _STATE["models"] 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) # Ensemble across all fold models by averaging logits (mean logit fusion). logits_sum = None with torch.inference_mode(): for m in models: m.to(device) lg = m(ids, masks).squeeze(0) logits_sum = lg if logits_sum is None else logits_sum + lg logits = (logits_sum / len(models)).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" num_folds = len(_STATE.get("models", [])) info_markdown = ( f"⏱ **{elapsed_ms:.0f} ms** response time · " f"Device: **{device_used}** · " f"Model: **ELECTRA-base discriminator (local {num_folds}-fold ensemble)** · " 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("""

🎯 Smart MCQ Solver

ELECTRA-base Neural Option Scorer  ·  MAP@3 = 0.8849 on 5-Fold CV

""") 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` | | **Checkpoint Files** | Local `results/electra_fold_{1..4}.pt` (full state_dict per fold, FP16) | | **Ensemble** | Mean logit fusion across all available folds | | **Scoring Head** | Joint 5-way scoring via 5× Dropout(0.2) + Linear layer | | **4-Fold Ensemble MAP@3** | **0.8849** (5-fold OOF reference) | | **4-Fold Ensemble Accuracy** | **~81%** (5-fold OOF reference) | | **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()