Update app.py

app.py

@@ -18,22 +18,22 @@ def solve(problem_text):
     if not problem_text or len(problem_text) < 10:
         return "// Error: Problem text too short."
 
-    # UPDATED PROMPT: Explicitly forbid comments
-    # We ask for a step-by-step algorithm to prevent one-liners, but still ban comments
-    prompt = f"Problem:\n{problem_text}\n\nWrite a complete, robust, step-by-step Python3 algorithm to solve this. DO NOT include any comments, docstrings, or conversational text. Only output the raw code.\nPython code solution:\n"
+    # THE FIX: Go back to the exact prompt that we know works!
+    prompt = f"Problem:\n{problem_text}\n\nPython code solution:\n"
     inputs = tokenizer(prompt, return_tensors="pt")
     
     with torch.no_grad():
         outputs = model.generate(
             input_ids=inputs["input_ids"],
             attention_mask=inputs["attention_mask"],
-            max_new_tokens=1024,  # REDUCED from 512 to speed up generation!
+            max_new_tokens=1024, 
             do_sample=False, 
             pad_token_id=tokenizer.eos_token_id
         )
     
     full_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
     
+    # 1. Extract the raw Python code block
     try:
         code_block = full_text.split("```python")[1]
         pure_code = code_block.split("```")[0]
@@ -42,14 +42,24 @@ def solve(problem_text):
             pure_code = full_text.split("Python code solution:\n")[1]
         else:
             pure_code = full_text
-            
-    # POST-PROCESSING: Scrub out any remaining comment lines
+
+    # 2. THE SCRUBBER: Delete the comments using Python, not the AI
     cleaned_lines = []
     for line in pure_code.split('\n'):
-        if not line.strip().startswith('#'):
-            cleaned_lines.append(line)
-            
-    return "\n".join(cleaned_lines).strip()
+        # This splits the line at the '#' symbol and only keeps the code on the left side
+        no_comment_line = line.split('#')[0].rstrip()
+        
+        # Only add the line if it's not completely empty
+        if no_comment_line.strip():
+            cleaned_lines.append(no_comment_line)
+
+    final_code = "\n".join(cleaned_lines).strip()
+    
+    # Safety net: If the scrubber accidentally deleted everything, return the original
+    if not final_code:
+        return pure_code.strip()
+        
+    return final_code
 
 demo = gr.Interface(fn=solve, inputs="text", outputs="text", api_name="predict")
 demo.launch()