Update app.py

app.py

@@ -12,30 +12,32 @@ import tempfile
 
 import tempfile
 
+import tempfile
+
 def review_code(code_snippet):
     print("✅ Received Code:", code_snippet)  # Debugging log
     
-    # ✅ Better instruction prompt
+    # ✅ Improved prompt with strict format rules
     prompt = f"""
 ### Instruction:
-You are a Python code reviewer. Your job is to analyze and fix errors in the provided Python code. 
+You are a Python code reviewer. Your job is to analyze and fix errors in the provided Python code.
 Make necessary corrections such as adding missing return statements, fixing syntax errors, and correcting logical mistakes. 
-Do NOT generate new functions or extra text—only return the fixed version of the provided code.
+ONLY return the corrected function definition—DO NOT generate any new function calls or explanations.
 
 ### Input Code:
 {code_snippet}
 
-### Reviewed Code:
+### Corrected Function:
 """
 
     # Process input
     inputs = tokenizer(prompt, return_tensors="pt").to("cpu")  # Move to CPU
     outputs = model.generate(
         **inputs,
-        max_new_tokens=50,  # ✅ Generate only 50 new tokens (fixes warning)
+        max_new_tokens=50,  # ✅ Limits AI output to prevent hallucinations
         do_sample=False,
-        num_beams=4,  # ✅ Ensures better correction quality
-        repetition_penalty=2.5  # ✅ Prevents repeated/unnecessary output
+        num_beams=4,  # ✅ Higher beams for better correction quality
+        repetition_penalty=3.0  # ✅ Stronger penalty to prevent repeated tokens
     )
 
     # Check if the model generated output
@@ -43,9 +45,13 @@ Do NOT generate new functions or extra text—only return the fixed version of t
         print("❌ Model did not generate output!")  # Debugging log
         return "Error: Model did not generate output."
 
-    reviewed_code = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    reviewed_code = tokenizer.decode(outputs[0], skip_special_tokens=True).strip()
     print("✅ Generated Code:", reviewed_code)  # Debugging log
 
+    # ✅ Ensure AI output starts with "def" to guarantee it returned a function
+    if not reviewed_code.startswith("def"):
+        return "Error: AI did not return a valid function."
+
     # ✅ Write reviewed code to a temporary file for download
     temp_file_path = tempfile.NamedTemporaryFile(delete=False, suffix=".txt").name
     with open(temp_file_path, "w") as temp_file: