Update app.py

app.py

@@ -18,28 +18,37 @@ def solve(problem_text):
     if not problem_text or len(problem_text) < 10:
         return "// Error: Problem text too short."
 
-    # Let's try a simpler, universal prompt format
     prompt = f"Problem:\n{problem_text}\n\nPython code solution:\n"
-    
     inputs = tokenizer(prompt, return_tensors="pt")
     
-    print("Starting generation...") # This will show up in HF Logs
-    
     with torch.no_grad():
         outputs = model.generate(
             input_ids=inputs["input_ids"],
             attention_mask=inputs["attention_mask"],
             max_new_tokens=512, 
-            do_sample=False, # Force deterministic greedy decoding
+            do_sample=False, 
             pad_token_id=tokenizer.eos_token_id
         )
     
     full_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
     
-    print("RAW MODEL OUTPUT:\n", full_text) # Check HF logs to see exactly what it did
-    
-    # TEMPORARY: Return the whole thing so you can see it in the GUI!
-    return full_text
+    # Extract only the Python code block
+    try:
+        # Split by the python code block marker
+        code_block = full_text.split("```python")[1]
+        
+        # Split again by the ending backticks and take everything before it
+        pure_code = code_block.split("```")[0]
+        
+        # Remove any leading or trailing blank lines
+        return pure_code.strip()
+        
+    except IndexError:
+        # Fallback: if the model forgot to use ```python formatting, 
+        # just return everything after "Python code solution:\n"
+        if "Python code solution:\n" in full_text:
+            return full_text.split("Python code solution:\n")[1].strip()
+        return full_text
 
 demo = gr.Interface(fn=solve, inputs="text", outputs="text", api_name="predict")
 demo.launch()