Upload app.py

app.py

@@ -1,7 +1,7 @@
-# app.py  ── Math-solution classifier for HF Spaces
-# Compatible with both LoRA-classification and LoRA-causal-LM adapters
-# Requirements (pin in requirements.txt):
-#   gradio torch transformers peft accelerate spaces
+# app.py  ── Math-solution classifier on HF Spaces (Zero-GPU-safe)
+#
+# Pin in requirements.txt:
+#   gradio==4.44.0 torch==2.1.0 transformers==4.35.0 peft==0.7.1 accelerate==0.25.0 spaces
 
 import os
 import json
@@ -9,107 +9,83 @@ import logging
 from typing import Tuple
 
 import gradio as gr
-import torch
-import spaces
-
-from transformers import (
-    AutoTokenizer,
-    AutoModelForSequenceClassification,
-)
-
-# PEFT imports (optional)
-try:
-    from peft.auto import (
-        AutoPeftModelForSequenceClassification,
-        AutoPeftModelForCausalLM,
-    )
-    PEFT_AVAILABLE = True
-except ImportError:  # PEFT not installed
-    PEFT_AVAILABLE = False
+import spaces  # <- Hugging Face Spaces SDK (Zero)
 
 # ──────────────────────────────────────────────────────────────────────────────
-# Config & logging
+# CONSTANTS (no CUDA use here)
 # ──────────────────────────────────────────────────────────────────────────────
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-
-ADAPTER_PATH   = os.getenv("ADAPTER_PATH", "./lora_adapter")  # local dir or Hub ID
+ADAPTER_PATH   = os.getenv("ADAPTER_PATH", "./lora_adapter")  # dir or Hub repo
 FALLBACK_MODEL = "distilbert-base-uncased"
 LABELS         = {0: "✅ Correct",
                   1: "🤔 Conceptual Error",
                   2: "🔢 Computational Error"}
 
-device = "cuda" if torch.cuda.is_available() else "cpu"
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
+# Globals that will live **inside the GPU worker**
 model     = None
 tokenizer = None
 model_ty  = None          # "classification" | "causal_lm" | "baseline"
 
 # ──────────────────────────────────────────────────────────────────────────────
-# Model loader
+# GPU-SIDE INITIALISATION & INFERENCE
 # ──────────────────────────────────────────────────────────────────────────────
-def load_model():
-    """Try adapter as classifier → causal-LM → plain baseline."""
+def _load_model_gpu():
+    """
+    Runs **inside the GPU worker**.
+    Tries LoRA classification adapter → LoRA causal-LM adapter → plain baseline.
+    """
     global model, tokenizer, model_ty
+    import torch
+    from transformers import (
+        AutoTokenizer,
+        AutoModelForSequenceClassification,
+    )
+    from peft.auto import (
+        AutoPeftModelForSequenceClassification,
+        AutoPeftModelForCausalLM,
+    )
 
-    dtype = torch.float16 if device == "cuda" else torch.float32
-
-    if PEFT_AVAILABLE and os.path.isdir(ADAPTER_PATH):
-        logger.info(f"Found adapter at {ADAPTER_PATH}")
+    dtype = torch.float16
+    if os.path.isdir(ADAPTER_PATH):
+        logger.info(f"[GPU] Loading adapter from {ADAPTER_PATH}")
 
-        # 1) Try sequence-classification adapter
-        try:
+        try:  # 1) classification adapter
             model = AutoPeftModelForSequenceClassification.from_pretrained(
-                ADAPTER_PATH,
-                torch_dtype=dtype,
-                device_map="auto" if device == "cuda" else None,
+                ADAPTER_PATH, torch_dtype=dtype, device_map="auto"
             )
             model_ty = "classification"
-            logger.info("Loaded adapter as sequence-classifier")
         except ValueError:
-            # 2) Fall back to causal-LM adapter
-            logger.info("Adapter is not a classifier – trying causal-LM")
+            logger.info("[GPU] Not a classifier, trying causal-LM")
             model = AutoPeftModelForCausalLM.from_pretrained(
-                ADAPTER_PATH,
-                torch_dtype=dtype,
-                device_map="auto" if device == "cuda" else None,
+                ADAPTER_PATH, torch_dtype=dtype, device_map="auto"
             )
             model_ty = "causal_lm"
-            logger.info("Loaded adapter as causal-LM")
 
         tokenizer = AutoTokenizer.from_pretrained(ADAPTER_PATH)
 
     else:
-        logger.warning("No adapter found – using baseline DistilBERT classifier")
+        logger.warning("[GPU] No adapter found – using baseline DistilBERT")
         tokenizer = AutoTokenizer.from_pretrained(FALLBACK_MODEL)
         model     = AutoModelForSequenceClassification.from_pretrained(
-            FALLBACK_MODEL,
-            num_labels=3,
-            ignore_mismatched_sizes=True,
-            torch_dtype=dtype,
+            FALLBACK_MODEL, num_labels=3, ignore_mismatched_sizes=True
         )
         model_ty = "baseline"
 
-    # Make sure we have a pad token
     if tokenizer.pad_token is None:
         tokenizer.pad_token = tokenizer.eos_token or tokenizer.sep_token
 
-    model.to(device)
     model.eval()
-    logger.info(f"Model ready on {device} as {model_ty}")
+    logger.info(f"[GPU] Model ready ({model_ty})")
+
 
-# ──────────────────────────────────────────────────────────────────────────────
-# Inference helpers
-# ──────────────────────────────────────────────────────────────────────────────
 def _classify_logits(question: str, solution: str) -> Tuple[str, str, str]:
+    import torch
     text = f"Question: {question}\n\nSolution:\n{solution}"
     inputs = tokenizer(
-        text,
-        return_tensors="pt",
-        padding=True,
-        truncation=True,
-        max_length=512,
-    ).to(device)
+        text, return_tensors="pt", padding=True, truncation=True, max_length=512
+    ).to("cuda")
 
     with torch.no_grad():
         logits = model(**inputs).logits
@@ -117,14 +93,14 @@ def _classify_logits(question: str, solution: str) -> Tuple[str, str, str]:
         pred   = int(torch.argmax(probs))
         conf   = f"{probs[pred].item():.3f}"
 
-    return LABELS.get(pred, "Unknown"), conf, "—"
+    return LABELS[pred], conf, "—"
+
 
 def _classify_generate(question: str, solution: str) -> Tuple[str, str, str]:
-    # Prompt must match the format you used in tuning
+    import torch
     prompt = (
         "You are a mathematics tutor.\n"
-        "You are given a math word problem and a student's solution. "
-        "Decide whether the solution is correct.\n\n"
+        "You are given a math word problem and a student's solution. Decide whether the solution is correct.\n\n"
         "- Correct            = all reasoning and calculations are correct.\n"
         "- Conceptual Error   = reasoning is wrong.\n"
         "- Computational Error= reasoning okay but arithmetic off.\n\n"
@@ -135,8 +111,7 @@ def _classify_generate(question: str, solution: str) -> Tuple[str, str, str]:
         f"Question: {question}\n\nSolution:\n{solution}\n\nAnswer:"
     )
 
-  
-    inputs = tokenizer(prompt, return_tensors="pt").to(device)
+    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
     with torch.no_grad():
         out_ids = model.generate(
             **inputs,
@@ -148,12 +123,9 @@ def _classify_generate(question: str, solution: str) -> Tuple[str, str, str]:
         skip_special_tokens=True,
     ).strip()
 
-
-    # Try to parse last JSON line
     verdict = "Unparsed"
     try:
-        line = generated.splitlines()[-1]
-        data = json.loads(line)
+        data = json.loads(generated.splitlines()[-1])
         v = data.get("verdict", "").lower()
         if v.startswith("corr"):
             verdict = LABELS[0]
@@ -166,21 +138,30 @@ def _classify_generate(question: str, solution: str) -> Tuple[str, str, str]:
 
     return verdict, "", generated
 
-def classify(question: str, solution: str):
+
+@spaces.GPU                         # <-- every CUDA op happens inside here
+def gpu_classify(question: str, solution: str):
+    """
+    Proxy target for Gradio. Executed in the GPU worker so CUDA is allowed.
+    Returns (verdict, confidence, raw_output)
+    """
+    if model is None:
+        _load_model_gpu()
+
     if not question.strip() or not solution.strip():
-        return "Please enter both fields.", "", ""
+        return "Please fill both fields.", "", ""
 
     if model_ty in ("classification", "baseline"):
         return _classify_logits(question, solution)
-    elif model_ty == "causal_lm":
+    else:  # causal_lm
         return _classify_generate(question, solution)
-    else:
-        return "Model not loaded.", "", ""
 
 # ──────────────────────────────────────────────────────────────────────────────
-# Build Gradio UI
+# CPU-SIDE UI (no torch.cuda here)
 # ──────────────────────────────────────────────────────────────────────────────
-load_model()
+def classify_proxy(q, s):
+    """Simple wrapper so Gradio can call the GPU function."""
+    return gpu_classify(q, s)
 
 with gr.Blocks(title="Math Solution Classifier") as demo:
     gr.Markdown("# 🧮 Math Solution Classifier")
@@ -191,15 +172,15 @@ with gr.Blocks(title="Math Solution Classifier") as demo:
 
     with gr.Row():
         with gr.Column():
-            q_in  = gr.Textbox(label="Math Question", lines=3)
-            s_in  = gr.Textbox(label="Proposed Solution", lines=6)
-            btn   = gr.Button("Classify", variant="primary")
+            q_in = gr.Textbox(label="Math Question", lines=3)
+            s_in = gr.Textbox(label="Proposed Solution", lines=6)
+            btn  = gr.Button("Classify", variant="primary")
         with gr.Column():
             verdict = gr.Textbox(label="Verdict", interactive=False)
             conf    = gr.Textbox(label="Confidence", interactive=False)
             raw     = gr.Textbox(label="Model Output", interactive=False)
 
-    btn.click(classify, [q_in, s_in], [verdict, conf, raw])
+    btn.click(classify_proxy, [q_in, s_in], [verdict, conf, raw])
 
     gr.Examples(
         [
@@ -210,11 +191,6 @@ with gr.Blocks(title="Math Solution Classifier") as demo:
         inputs=[q_in, s_in],
     )
 
-
-@spaces.GPU                       # or @spaces.CPU if you deploy on CPU
-
+@spaces.CPU                         # UI served from the CPU worker
 def launch_app():
-    return demo                   # the Gradio Blocks object you built
-
-if __name__ == "__main__":
-    demo.launch()
+    return demo