Upload app.py

app.py

@@ -19,11 +19,10 @@ def load_model():
     global model, tokenizer
     
     try:
-        from peft import AutoPeftModelForSequenceClassification
+        from peft import AutoPeftModelForCausalLM  # Changed from SequenceClassification
         
-        # Load the LoRA adapter model
-        # The adapter files should be in a folder (e.g., "./lora_adapter")
-        model = AutoPeftModelForSequenceClassification.from_pretrained(
+        # Load the LoRA adapter model for text generation
+        model = AutoPeftModelForCausalLM.from_pretrained(
             "./lora_adapter",  # Path to your adapter files
             torch_dtype=torch.float16,
             device_map="auto"
@@ -52,28 +51,58 @@ def load_model():
         if tokenizer.pad_token is None:
             tokenizer.pad_token = tokenizer.eos_token
             
-        model = AutoModelForSequenceClassification.from_pretrained(
-            model_name, 
-            num_labels=3,
-            ignore_mismatched_sizes=True
-        )
+        from transformers import AutoModelForCausalLM
+        model = AutoModelForCausalLM.from_pretrained(model_name)
         
         return f"Fallback model loaded. LoRA error: {e}"
 
+def get_system_prompt():
+    """Generates the specific system prompt for the fine-tuning task."""
+    return """You are a mathematics tutor.
+You are given a math word problem, and a solution written by a student.
+Analyze the solution carefully, line-by-line, and classify it into one of the following categories:
+- Correct (All logic is correct, and all calculations are correct)
+- Conceptual Error (There is an error in reasoning or logic somewhere in the solution)
+- Computational Error (All logic and reasoning is correct, but the result of some calculation is incorrect)
+Respond *only* with a valid JSON object that follows this exact schema:
+```json
+{
+  "verdict": "must be one of 'correct', 'conceptual_error', or 'computational_error'",
+  "erroneous_line": "the exact, verbatim text of the first incorrect line, or null if the verdict is 'correct'",
+  "explanation": "a brief, one-sentence explanation of the error, or null if the verdict is 'correct'"
+}
+```
+Do NOT add any text or explanations before or after the JSON object.
+"""
+
 def classify_solution(question: str, solution: str):
     """
-    Classify the math solution
+    Classify the math solution using the exact training format
     Returns: (classification_label, confidence_score, explanation)
     """
     if not question.strip() or not solution.strip():
-        return "Please fill in both fields", 0.0, ""
+        return "Please fill in both fields", "", ""
     
     if not model or not tokenizer:
-        return "Model not loaded", 0.0, ""
+        return "Model not loaded", "", ""
     
     try:
-        # Combine question and solution for input
-        text_input = f"Question: {question}\nSolution: {solution}"
+        # Create the exact prompt format used in training
+        system_prompt = get_system_prompt()
+        user_message = f"Problem: {question}\n\nSolution:\n{solution}"
+        
+        # Format as chat messages (common for instruction-tuned models)
+        messages = [
+            {"role": "system", "content": system_prompt},
+            {"role": "user", "content": user_message}
+        ]
+        
+        # Apply chat template
+        text_input = tokenizer.apply_chat_template(
+            messages, 
+            tokenize=False, 
+            add_generation_token=True
+        )
         
         # Tokenize input
         inputs = tokenizer(
@@ -81,32 +110,51 @@ def classify_solution(question: str, solution: str):
             return_tensors="pt",
             truncation=True,
             padding=True,
-            max_length=512
+            max_length=2048  # Increased for longer prompts
         )
         
-        # Get model prediction
+        # Generate response (not just classify)
         with torch.no_grad():
-            outputs = model(**inputs)
-            predictions = torch.nn.functional.softmax(outputs.logits, dim=-1)
-            predicted_class = torch.argmax(predictions, dim=-1).item()
-            confidence = predictions[0][predicted_class].item()
-        
-        classification = label_mapping[predicted_class]
+            outputs = model.generate(
+                **inputs,
+                max_new_tokens=200,
+                temperature=0.1,
+                do_sample=True,
+                pad_token_id=tokenizer.pad_token_id
+            )
         
-        # Create explanation based on classification
-        explanations = {
-            0: "The mathematical approach and calculations are both sound.",
-            1: "The approach or understanding has fundamental issues.",
-            2: "The approach is correct, but there are calculation errors."
-        }
+        # Decode the generated response
+        generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
         
-        explanation = explanations[predicted_class]
+        # Extract just the JSON response (after the input)
+        response_start = generated_text.find(text_input) + len(text_input)
+        json_response = generated_text[response_start:].strip()
         
-        return classification, f"{confidence:.2%}", explanation
+        # Parse the JSON response
+        import json
+        try:
+            result = json.loads(json_response)
+            verdict = result.get("verdict", "unknown")
+            erroneous_line = result.get("erroneous_line", "")
+            explanation = result.get("explanation", "")
+            
+            # Map verdict to display format
+            verdict_mapping = {
+                "correct": "✅ Correct",
+                "conceptual_error": "🤔 Conceptual Error", 
+                "computational_error": "🔢 Computational Error"
+            }
+            
+            display_verdict = verdict_mapping.get(verdict, f"❓ {verdict}")
+            
+            return display_verdict, erroneous_line or "None", explanation or "Solution is correct"
+            
+        except json.JSONDecodeError:
+            return f"Model response: {json_response}", "", "Could not parse JSON response"
         
     except Exception as e:
         logger.error(f"Error during classification: {e}")
-        return f"Classification error: {str(e)}", "0%", ""
+        return f"Classification error: {str(e)}", "", ""
 
 # Load model on startup
 load_model()
@@ -134,7 +182,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)
-            confidence_output = gr.Textbox(label="Confidence", interactive=False)
+            erroneous_line_output = gr.Textbox(label="Erroneous Line", interactive=False)
             explanation_output = gr.Textbox(label="Explanation", interactive=False, lines=3)
     
     # Examples
@@ -159,7 +207,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, confidence_output, explanation_output]
+        outputs=[classification_output, erroneous_line_output, explanation_output]
     )
 
 if __name__ == "__main__":