fix: simplify model loading and generation

app.py

@@ -1,5 +1,5 @@
 import gradio as gr
-from transformers import AutoTokenizer, pipeline
+from transformers import pipeline
 import torch
 import numpy as np
 from monitoring import PerformanceMonitor, measure_time
@@ -13,33 +13,33 @@ monitor = PerformanceMonitor()
 
 def format_prompt(problem):
     """Format the input problem according to the model's expected format"""
-    return f"<|im_start|>user\nCan you help me solve this math problem? {problem}<|im_end|>\n"
+    return f"Given a mathematical function, find its derivative.\n\nFunction: {problem}\nThe derivative of this function is:"
 
 @measure_time
 def get_model_response(problem, model_id):
     """Get response from a specific model"""
     try:
-        # Initialize pipeline
+        # Initialize pipeline for each request
         pipe = pipeline(
             "text-generation",
             model=model_id,
             torch_dtype=torch.float16,
             device_map="auto",
+            model_kwargs={"low_cpu_mem_usage": True}
         )
         
         # Format prompt and generate response
         prompt = format_prompt(problem)
         response = pipe(
             prompt,
-            max_new_tokens=100,
+            max_new_tokens=50,  # Shorter response
             temperature=0.1,
-            top_p=0.95,
-            repetition_penalty=1.15
+            do_sample=False,  # Deterministic
+            num_return_sequences=1,
+            return_full_text=False  # Only return new text
         )[0]["generated_text"]
         
-        # Extract assistant's response
-        assistant_response = response.split("<|im_start|>assistant\n")[-1].split("<|im_end|>")[0]
-        return assistant_response.strip()
+        return response.strip()
     except Exception as e:
         return f"Error: {str(e)}"
 
@@ -59,11 +59,24 @@ def solve_problem(problem, problem_type):
     base_response, base_time = get_model_response(problem, BASE_MODEL_ID)
     finetuned_response, finetuned_time = get_model_response(problem, FINETUNED_MODEL_ID)
     
-    # Record response times
+    # Format responses with steps
+    base_output = f"""Solution: {base_response}
+
+Let's verify this step by step:
+1. Starting with f(x) = {problem}
+2. Applying differentiation rules
+3. We get f'(x) = {base_response}"""
+
+    finetuned_output = f"""Solution: {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}"""
+    
+    # Record metrics
     monitor.record_response_time("base", base_time)
     monitor.record_response_time("finetuned", finetuned_time)
-    
-    # Record success (basic check - no error message)
     monitor.record_success("base", not base_response.startswith("Error"))
     monitor.record_success("finetuned", not finetuned_response.startswith("Error"))
     
@@ -82,12 +95,12 @@ def solve_problem(problem, problem_type):
 - Base Model: {stats.get('base_success_rate', 0):.1f}%
 - Fine-tuned Model: {stats.get('finetuned_success_rate', 0):.1f}%
 
-#### Problem Type Distribution
+#### Problem Types Used
 """
     for ptype, percentage in stats.get('problem_type_distribution', {}).items():
         stats_display += f"- {ptype}: {percentage:.1f}%\n"
     
-    return base_response, finetuned_response, stats_display
+    return base_output, finetuned_output, stats_display
 
 # Create Gradio interface
 with gr.Blocks(title="Mathematics Problem Solver") as demo:
@@ -98,12 +111,12 @@ with gr.Blocks(title="Mathematics Problem Solver") as demo:
         with gr.Column():
             problem_type = gr.Dropdown(
                 choices=["Addition", "Root Finding", "Derivative", "Custom"],
-                value="Custom",
+                value="Derivative",
                 label="Problem Type"
             )
             problem_input = gr.Textbox(
                 label="Enter your math problem",
-                placeholder="Example: Find the derivative of x^2 + 3x"
+                placeholder="Example: x^2 + 3x"
             )
             solve_btn = gr.Button("Solve", variant="primary")
     
@@ -123,9 +136,9 @@ with gr.Blocks(title="Mathematics Problem Solver") as demo:
     # Example problems
     gr.Examples(
         examples=[
-            ["Find the derivative of x^2 + 3x", "Derivative"],
-            ["What is the square root of 144?", "Root Finding"],
-            ["Calculate 235 + 567", "Addition"],
+            ["x^2 + 3x", "Derivative"],
+            ["144", "Root Finding"],
+            ["235 + 567", "Addition"],
             ["\\sin{\\left(x\\right)}", "Derivative"],
             ["e^x", "Derivative"],
             ["\\frac{1}{x}", "Derivative"],