Update app.py

app.py

@@ -6,10 +6,14 @@ import gradio as gr
 
 # Ensure the environment has access to a CUDA-capable GPU
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print(f"Using device: {device}")
 
 # Load model and tokenizer directly to GPU if available
+print("Loading tokenizer...")
 tokenizer = AutoTokenizer.from_pretrained("microsoft/Phi-3-mini-4k-instruct", trust_remote_code=True)
-model = AutoModelForCausalLM.from_pretrained("microsoft/Phi-3-mini-4k-instruct", trust_remote_code=True, device_map="auto")
+
+print("Loading model...")
+model = AutoModelForCausalLM.from_pretrained("microsoft/Phi-3-mini-4k-instruct", trust_remote_code=True, device_map="auto").to(device)
 
 # Define templates for problems
 templates = {
@@ -52,21 +56,24 @@ def generate_synthetic_math_problems(num_problems):
     return problems
 
 def solve_problem(problem):
-    # Encode the problem
-    inputs = tokenizer(problem, return_tensors="pt").to(device)
-    
-    # Generate a response from the model
-    outputs = model.generate(inputs["input_ids"], max_length=100)
-    
-    # Decode the response
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    
-    # Strip the answer to only the math (assuming answer is preceded by "The answer is ")
-    if "The answer is " in response:
-        answer = response.split("The answer is ")[-1].strip()
-    else:
-        answer = response.strip()
+    print(f"Solving problem: {problem}")
+    with torch.no_grad():
+        # Encode the problem
+        inputs = tokenizer(problem, return_tensors="pt").to(device)
+        
+        # Generate a response from the model
+        outputs = model.generate(inputs["input_ids"], max_length=50, num_return_sequences=1, do_sample=True)
+        
+        # Decode the response
+        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        
+        # Strip the answer to only the math (assuming answer is preceded by "The answer is ")
+        if "The answer is " in response:
+            answer = response.split("The answer is ")[-1].strip()
+        else:
+            answer = response.strip()
     
+    print(f"Problem: {problem}, Answer: {answer}")
     return answer
 
 def generate_and_solve_problems(num_problems):
@@ -83,6 +90,7 @@ def generate_and_solve_problems(num_problems):
     return solved_problems
 
 def gradio_interface(num_problems):
+    print(f"Generating and solving {num_problems} problems...")
     solved_problems = generate_and_solve_problems(num_problems)
     return json.dumps(solved_problems, indent=4)