Upload app.py

app.py

@@ -1,18 +1,28 @@
 # app.py  ── Math-solution classifier for HF Spaces
-# Requires: gradio, torch, transformers, peft, accelerate, spaces
+# Compatible with both LoRA-classification and LoRA-causal-LM adapters
+# Requirements (pin in requirements.txt):
+#   gradio torch transformers peft accelerate spaces
 
 import os
+import json
 import logging
+from typing import Tuple
 
 import gradio as gr
 import torch
-from transformers import AutoTokenizer, AutoModelForSequenceClassification
+from transformers import (
+    AutoTokenizer,
+    AutoModelForSequenceClassification,
+)
 
-# Optional PEFT import (only available if you include it in requirements.txt)
+# PEFT imports (optional)
 try:
-    from peft import AutoPeftModelForSequenceClassification
+    from peft.auto import (
+        AutoPeftModelForSequenceClassification,
+        AutoPeftModelForCausalLM,
+    )
     PEFT_AVAILABLE = True
-except ImportError:
+except ImportError:  # PEFT not installed
     PEFT_AVAILABLE = False
 
 # ──────────────────────────────────────────────────────────────────────────────
@@ -21,55 +31,75 @@ except ImportError:
 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")  # local dir or Hub ID
 FALLBACK_MODEL = "distilbert-base-uncased"
-LABELS = {0: "✅ Correct",
-          1: "🤔 Conceptual Error",
-          2: "🔢 Computational Error"}
+LABELS         = {0: "✅ Correct",
+                  1: "🤔 Conceptual Error",
+                  2: "🔢 Computational Error"}
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
-model = None
+
+model     = None
 tokenizer = None
+model_ty  = None          # "classification" | "causal_lm" | "baseline"
 
 # ──────────────────────────────────────────────────────────────────────────────
-# Load model & tokenizer
+# Model loader
 # ──────────────────────────────────────────────────────────────────────────────
 def load_model():
-    """Load the LoRA adapter if present, otherwise a baseline classifier."""
-    global model, tokenizer
+    """Try adapter as classifier → causal-LM → plain baseline."""
+    global model, tokenizer, model_ty
+
+    dtype = torch.float16 if device == "cuda" else torch.float32
 
     if PEFT_AVAILABLE and os.path.isdir(ADAPTER_PATH):
-        logger.info(f"Loading LoRA adapter from {ADAPTER_PATH}")
-        model = AutoPeftModelForSequenceClassification.from_pretrained(
-            ADAPTER_PATH,
-            torch_dtype=torch.float16 if device == "cuda" else torch.float32,
-            device_map="auto" if device == "cuda" else None,
-        )
+        logger.info(f"Found adapter at {ADAPTER_PATH}")
+
+        # 1) Try sequence-classification adapter
+        try:
+            model = AutoPeftModelForSequenceClassification.from_pretrained(
+                ADAPTER_PATH,
+                torch_dtype=dtype,
+                device_map="auto" if device == "cuda" else None,
+            )
+            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")
+            model = AutoPeftModelForCausalLM.from_pretrained(
+                ADAPTER_PATH,
+                torch_dtype=dtype,
+                device_map="auto" if device == "cuda" else None,
+            )
+            model_ty = "causal_lm"
+            logger.info("Loaded adapter as causal-LM")
+
         tokenizer = AutoTokenizer.from_pretrained(ADAPTER_PATH)
+
     else:
-        logger.warning("LoRA adapter not found – falling back to baseline model")
+        logger.warning("No adapter found – using baseline DistilBERT classifier")
         tokenizer = AutoTokenizer.from_pretrained(FALLBACK_MODEL)
-        model = AutoModelForSequenceClassification.from_pretrained(
+        model     = AutoModelForSequenceClassification.from_pretrained(
             FALLBACK_MODEL,
             num_labels=3,
             ignore_mismatched_sizes=True,
+            torch_dtype=dtype,
         )
+        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("Model & tokenizer ready")
+    logger.info(f"Model ready on {device} as {model_ty}")
 
 # ─────────────────────────���────────────────────────────────────────────────────
-# Inference helper
+# Inference helpers
 # ──────────────────────────────────────────────────────────────────────────────
-def classify(question: str, solution: str):
-    """Return (label, confidence, placeholder-explanation)."""
-    if not question.strip() or not solution.strip():
-        return "Please provide both question and solution.", "", ""
-
+def _classify_logits(question: str, solution: str) -> Tuple[str, str, str]:
     text = f"Question: {question}\n\nSolution:\n{solution}"
     inputs = tokenizer(
         text,
@@ -81,11 +111,64 @@ def classify(question: str, solution: str):
 
     with torch.no_grad():
         logits = model(**inputs).logits
-        probs = torch.softmax(logits, dim=-1)[0]
-        pred = int(torch.argmax(probs))
-        confidence = f"{probs[pred].item():.3f}"
+        probs  = torch.softmax(logits, dim=-1)[0]
+        pred   = int(torch.argmax(probs))
+        conf   = f"{probs[pred].item():.3f}"
+
+    return LABELS.get(pred, "Unknown"), conf, "—"
+
+def _classify_generate(question: str, solution: str) -> Tuple[str, str, str]:
+    # Prompt must match the format you used in tuning
+    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"
+        "- Correct  = all reasoning and calculations are correct.\n"
+        "- Conceptual Error = reasoning is wrong.\n"
+        "- Computational Error = reasoning okay but arithmetic off.\n\n"
+        "Reply with ONLY one of these JSON lines:\n"
+        '{"verdict": "correct"}\n'
+        '{"verdict": "conceptual"}\n'
+        '{"verdict": "computational"}\n\n"
+        f"Question: {question}\n\nSolution:\n{solution}\n\nAnswer:"
+    )
+
+    inputs = tokenizer(prompt, return_tensors="pt").to(device)
+    with torch.no_grad():
+        out_ids = model.generate(
+            **inputs,
+            max_new_tokens=32,
+            pad_token_id=tokenizer.eos_token_id,
+        )
+    generated = tokenizer.decode(out_ids[0][inputs["input_ids"].shape[1]:],
+                                 skip_special_tokens=True).strip()
+
+    # Try to parse last JSON line
+    verdict = "Unparsed"
+    try:
+        line = generated.splitlines()[-1]
+        data = json.loads(line)
+        v = data.get("verdict", "").lower()
+        if v.startswith("corr"):
+            verdict = LABELS[0]
+        elif v.startswith("conc"):
+            verdict = LABELS[1]
+        elif v.startswith("comp"):
+            verdict = LABELS[2]
+    except Exception:
+        pass
+
+    return verdict, "", generated
 
-    return LABELS.get(pred, "Unknown"), confidence, "—"
+def classify(question: str, solution: str):
+    if not question.strip() or not solution.strip():
+        return "Please enter both fields.", "", ""
+
+    if model_ty in ("classification", "baseline"):
+        return _classify_logits(question, solution)
+    elif model_ty == "causal_lm":
+        return _classify_generate(question, solution)
+    else:
+        return "Model not loaded.", "", ""
 
 # ──────────────────────────────────────────────────────────────────────────────
 # Build Gradio UI
@@ -95,27 +178,27 @@ load_model()
 with gr.Blocks(title="Math Solution Classifier") as demo:
     gr.Markdown("# 🧮 Math Solution Classifier")
     gr.Markdown(
-        "Classify a student’s math solution as **correct**, **conceptually flawed**, "
+        "Classify a student's math solution as **correct**, **conceptually flawed**, "
         "or **computationally flawed**."
     )
 
     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)
-            expl = gr.Textbox(label="Explanation", 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, expl])
+    btn.click(classify, [q_in, s_in], [verdict, conf, raw])
 
     gr.Examples(
         [
             ["Solve for x: 2x + 5 = 13", "2x + 5 = 13\n2x = 8\nx = 4"],
             ["Find the derivative of f(x)=x²", "f'(x)=2x+1"],
-            ["What is 15 % of 200?", "0.15 × 200 = 30"],
+            ["What is 15% of 200?", "0.15 × 200 = 30"],
         ],
         inputs=[q_in, s_in],
     )