Upload app.py

app.py

@@ -15,18 +15,17 @@ tokenizer = None
 label_mapping = {0: "✅ Correct", 1: "🤔 Conceptually Flawed", 2: "🔢 Computationally Flawed"}
 
 def load_model():
-    """Load your trained LoRA adapter with base model"""
+    """Load your trained LoRA adapter with classification head"""
     global model, tokenizer
     
     try:
-        from peft import AutoPeftModelForCausalLM  # Changed from SequenceClassification
+        from peft import AutoPeftModelForSequenceClassification  # Back to classification
         
-        # Load the LoRA adapter model for text generation
-        model = AutoPeftModelForCausalLM.from_pretrained(
+        # Load the LoRA adapter model for classification
+        model = AutoPeftModelForSequenceClassification.from_pretrained(
             "./lora_adapter",  # Path to your adapter files
-            torch_dtype=torch.float32,  # Use float32 for CPU
-            device_map="cpu",  # Force CPU
-            low_cpu_mem_usage=True  # Optimize for low memory
+            torch_dtype=torch.float16,
+            device_map="auto"
         )
         
         # Load tokenizer from the same directory
@@ -37,23 +36,27 @@ def load_model():
             tokenizer.pad_token = tokenizer.eos_token
             logger.info("Set pad_token to eos_token")
         
-        logger.info("LoRA model loaded successfully")
-        return "LoRA model loaded successfully!"
+        logger.info("LoRA classification model loaded successfully")
+        return "LoRA classification model loaded successfully!"
         
     except Exception as e:
         logger.error(f"Error loading LoRA model: {e}")
         # Fallback to placeholder for testing
         logger.warning("Using placeholder model loading - replace with your actual model!")
         
-        model_name = "microsoft/DialoGPT-medium"  # Closer to Phi-4 architecture
+        model_name = "distilbert-base-uncased"  # Simple fallback
         tokenizer = AutoTokenizer.from_pretrained(model_name)
         
         # Fix padding token for fallback model too
         if tokenizer.pad_token is None:
             tokenizer.pad_token = tokenizer.eos_token
             
-        from transformers import AutoModelForCausalLM
-        model = AutoModelForCausalLM.from_pretrained(model_name)
+        from transformers import AutoModelForSequenceClassification
+        model = AutoModelForSequenceClassification.from_pretrained(
+            model_name, 
+            num_labels=3,
+            ignore_mismatched_sizes=True
+        )
         
         return f"Fallback model loaded. LoRA error: {e}"
 
@@ -189,7 +192,7 @@ with gr.Blocks(title="Math Solution Classifier", theme=gr.themes.Soft()) as app:
         
         with gr.Column():
             classification_output = gr.Textbox(label="Classification", interactive=False)
-            erroneous_line_output = gr.Textbox(label="Erroneous Line", interactive=False)
+            confidence_output = gr.Textbox(label="Confidence", interactive=False)
             explanation_output = gr.Textbox(label="Explanation", interactive=False, lines=3)
     
     # Examples
@@ -214,7 +217,7 @@ with gr.Blocks(title="Math Solution Classifier", theme=gr.themes.Soft()) as app:
     classify_btn.click(
         fn=classify_solution,
         inputs=[question_input, solution_input],
-        outputs=[classification_output, erroneous_line_output, explanation_output]
+        outputs=[classification_output, confidence_output, explanation_output]
     )
 
 if __name__ == "__main__":