Update app.py

app.py

@@ -18,37 +18,37 @@ def solve(problem_text):
     if not problem_text or len(problem_text) < 10:
         return "// Error: Problem text too short."
 
-    prompt = f"Problem:\n{problem_text}\n\nPython code solution:\n"
+    # UPDATED PROMPT: Explicitly forbid comments
+    prompt = f"Problem:\n{problem_text}\n\nProvide ONLY the Python3 code. DO NOT include any comments, explanations, or docstrings.\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=512, 
+            max_new_tokens=300,  # REDUCED from 512 to speed up generation!
             do_sample=False, 
             pad_token_id=tokenizer.eos_token_id
         )
     
     full_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
     
-    # 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
+            pure_code = full_text.split("Python code solution:\n")[1]
+        else:
+            pure_code = full_text
+            
+    # POST-PROCESSING: Scrub out any remaining comment lines
+    cleaned_lines = []
+    for line in pure_code.split('\n'):
+        if not line.strip().startswith('#'):
+            cleaned_lines.append(line)
+            
+    return "\n".join(cleaned_lines).strip()
 
 demo = gr.Interface(fn=solve, inputs="text", outputs="text", api_name="predict")
 demo.launch()