updated model.py

model.py

@@ -22,14 +22,14 @@ model = AutoModelForCausalLM.from_pretrained(
 ).to(device)
 
 def generate_code(prompt: str, max_tokens: int = 256):
-    formatted_prompt = f"# Python\n{prompt}\n\n"  # Ensure the model understands it's code
+    formatted_prompt = f"{prompt}\n### Code:\n"  # Ensure the model understands it's code
 
     inputs = tokenizer(
         formatted_prompt, 
         return_tensors="pt", 
         padding=True, 
         truncation=True,  
-        max_length=1024  # Explicit max length to prevent issues
+        max_length=512  # Explicit max length to prevent issues
     ).to(device)
 
     output = model.generate(