Spaces:
Running on Zero
Running on Zero
| """ | |
| 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: | |
| def run_electra_inference(prompt, options): | |
| return run_electra_core(prompt, options) | |
| else: | |
| 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(""" | |
| <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 Neural Option Scorer · | |
| <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` | | |
| | **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() | |