feat: add model comparison with base and fine-tuned

app.py

@@ -17,7 +17,7 @@ tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
 tokenizer.pad_token = tokenizer.eos_token
 
 print("Loading base model...")
-model = AutoModelForCausalLM.from_pretrained(
+base_model = AutoModelForCausalLM.from_pretrained(
     BASE_MODEL,
     device_map="auto",
     torch_dtype=torch.float16,
@@ -25,14 +25,17 @@ model = AutoModelForCausalLM.from_pretrained(
     use_safetensors=True
 )
 
-print("Loading LoRA adapter...")
-model = PeftModel.from_pretrained(
-    model, 
+print("Loading fine-tuned model...")
+finetuned_model = PeftModel.from_pretrained(
+    base_model, 
     ADAPTER_MODEL,
     torch_dtype=torch.float16,
     device_map="auto"
 )
-model.eval()
+
+# Set models to eval mode
+base_model.eval()
+finetuned_model.eval()
 
 def format_prompt(problem: str, problem_type: str) -> str:
     """Format input prompt for the model"""
@@ -46,7 +49,7 @@ The derivative of this function is:"""
 
 Problem: {problem}
 The solution is:"""
-    else:  # Roots or Custom
+    else:  # Roots
         return f"""Find the roots of this equation.
 
 Equation: {problem}
@@ -54,7 +57,7 @@ The roots are:"""
 
 @spaces.GPU
 @measure_time
-def get_model_response(problem: str, problem_type: str) -> str:
+def get_model_response(problem: str, problem_type: str, model) -> str:
     """Generate response from model"""
     # Format prompt
     prompt = format_prompt(problem, problem_type)
@@ -81,42 +84,68 @@ def get_model_response(problem: str, problem_type: str) -> str:
 
 @spaces.GPU
 def solve_problem(problem: str, problem_type: str) -> tuple:
-    """Solve math problem and track performance"""
+    """Solve math problem with both models"""
     if not problem:
-        return "Please enter a problem", None
+        return "Please enter a problem", "Please enter a problem", None
     
     # Record problem type
     monitor.record_problem_type(problem_type)
     
-    # Get model response with timing
-    response, time_taken = get_model_response(problem, problem_type)
+    # Get responses from both models with timing
+    base_response, base_time = get_model_response(problem, problem_type, base_model)
+    finetuned_response, finetuned_time = get_model_response(problem, problem_type, finetuned_model)
     
-    # Format output with steps
+    # Format outputs with steps
     if problem_type == "Derivative":
-        output = f"""Generated derivative: {response}
+        base_output = f"""Generated derivative: {base_response}
 
 Let's verify this step by step:
 1. Starting with f(x) = {problem}
 2. Applying differentiation rules
-3. We get f'(x) = {response}"""
+3. We get f'(x) = {base_response}"""
+
+        finetuned_output = f"""Generated derivative: {finetuned_response}
+
+Let's verify this step by step:
+1. Starting with f(x) = {problem}
+2. Applying differentiation rules
+3. We get f'(x) = {finetuned_response}"""
+
     elif problem_type == "Addition":
-        output = f"""Solution: {response}
+        base_output = f"""Solution: {base_response}
+
+Let's verify this step by step:
+1. Starting with: {problem}
+2. Adding the numbers
+3. We get: {base_response}"""
+
+        finetuned_output = f"""Solution: {finetuned_response}
 
 Let's verify this step by step:
 1. Starting with: {problem}
 2. Adding the numbers
-3. We get: {response}"""
+3. We get: {finetuned_response}"""
+
     else:  # Roots
-        output = f"""Found roots: {response}
+        base_output = f"""Found roots: {base_response}
 
 Let's verify this step by step:
 1. Starting with equation: {problem}
 2. Solving for x
-3. Roots are: {response}"""
+3. Roots are: {base_response}"""
+
+        finetuned_output = f"""Found roots: {finetuned_response}
+
+Let's verify this step by step:
+1. Starting with equation: {problem}
+2. Solving for x
+3. Roots are: {finetuned_response}"""
     
     # Record metrics
-    monitor.record_response_time("model", time_taken)
-    monitor.record_success("model", not response.startswith("Error"))
+    monitor.record_response_time("base", base_time)
+    monitor.record_response_time("finetuned", finetuned_time)
+    monitor.record_success("base", not base_response.startswith("Error"))
+    monitor.record_success("finetuned", not finetuned_response.startswith("Error"))
     
     # Get updated statistics
     stats = monitor.get_statistics()
@@ -125,23 +154,25 @@ Let's verify this step by step:
     stats_display = f"""
 ### Performance Metrics
 
-#### Response Times
-- Average: {stats.get('model_avg_response_time', 0):.2f} seconds
+#### Response Times (seconds)
+- Base Model: {stats.get('base_avg_response_time', 0):.2f} avg
+- Fine-tuned Model: {stats.get('finetuned_avg_response_time', 0):.2f} avg
 
-#### Success Rate
-- {stats.get('model_success_rate', 0):.1f}%
+#### Success Rates
+- Base Model: {stats.get('base_success_rate', 0):.1f}%
+- Fine-tuned Model: {stats.get('finetuned_success_rate', 0):.1f}%
 
 #### Problem Types Used
 """
     for ptype, percentage in stats.get('problem_type_distribution', {}).items():
         stats_display += f"- {ptype}: {percentage:.1f}%\n"
     
-    return output, stats_display
+    return base_output, finetuned_output, stats_display
 
 # Create Gradio interface
 with gr.Blocks(title="Mathematics Problem Solver") as demo:
     gr.Markdown("# Mathematics Problem Solver")
-    gr.Markdown("Using our fine-tuned model to solve mathematical problems")
+    gr.Markdown("Compare solutions between base and fine-tuned models")
     
     with gr.Row():
         with gr.Column():
@@ -157,10 +188,13 @@ with gr.Blocks(title="Mathematics Problem Solver") as demo:
             solve_btn = gr.Button("Solve", variant="primary")
     
     with gr.Row():
-        solution_output = gr.Textbox(
-            label="Solution with Steps",
-            lines=6
-        )
+        with gr.Column():
+            gr.Markdown("### Base Model")
+            base_output = gr.Textbox(label="Base Model Solution", lines=6)
+        
+        with gr.Column():
+            gr.Markdown("### Fine-tuned Model")
+            finetuned_output = gr.Textbox(label="Fine-tuned Model Solution", lines=6)
     
     # Performance metrics display
     with gr.Row():
@@ -177,7 +211,7 @@ with gr.Blocks(title="Mathematics Problem Solver") as demo:
             ["\\frac{1}{x}", "Derivative"]
         ],
         inputs=[problem_input, problem_type],
-        outputs=[solution_output, metrics_display],
+        outputs=[base_output, finetuned_output, metrics_display],
         fn=solve_problem,
         cache_examples=False  # Disable caching
     )
@@ -186,7 +220,7 @@ with gr.Blocks(title="Mathematics Problem Solver") as demo:
     solve_btn.click(
         fn=solve_problem,
         inputs=[problem_input, problem_type],
-        outputs=[solution_output, metrics_display]
+        outputs=[base_output, finetuned_output, metrics_display]
     )
 
 if __name__ == "__main__":