udpated

main.py

@@ -80,17 +80,15 @@ def run_ai_inference(diff: str) -> str:
     if not model or not tokenizer:
         raise RuntimeError("Model is not loaded.")
 
-    # Improved prompt for codegen-350M-mono
-    prompt = (
-        "Below is a Python function. Please provide a code review comment with suggestions for improvement, in natural language. "
-        "Do not repeat the code.\n"
-        f"{diff[:800]}\n"
-        "Review comment:"
-    )
+    # Simple, direct prompt for codegen-350M-mono
+    prompt = f"""Code:
+{diff[:500]}
+
+Review: This code could be improved by adding"""
     encoded = tokenizer(
         prompt,
         return_tensors="pt",
-        max_length=1024,
+        max_length=512,  # Reduced from 1024 for faster processing
         truncation=True,
         padding="max_length"
     )
@@ -100,21 +98,38 @@ def run_ai_inference(diff: str) -> str:
         outputs = model.generate(
             input_ids=input_ids,
             attention_mask=attention_mask,
-            max_new_tokens=128,
+            max_new_tokens=32,  # Further reduced for speed
             do_sample=True,
-            temperature=0.7,
-            top_p=0.95,
+            temperature=0.9,
+            top_p=0.85,
             num_return_sequences=1,
             pad_token_id=tokenizer.eos_token_id if tokenizer.eos_token_id is not None else tokenizer.pad_token_id,
             eos_token_id=tokenizer.eos_token_id if tokenizer.eos_token_id is not None else tokenizer.pad_token_id,
             use_cache=True
         )
     response_text = tokenizer.decode(outputs[0][input_ids.shape[1]:], skip_special_tokens=True)
-    # Post-process: filter out code-like lines and fallback if needed
-    review_lines = [line.strip() for line in response_text.strip().split('\n') if line.strip()]
-    # Filter out lines that look like code
-    comment_lines = [l for l in review_lines if not l.startswith("def ") and not l.startswith("class ") and not l.endswith(":") and not l.startswith("#")]
-    review = comment_lines[0] if comment_lines else "Consider adding a docstring and input validation."
+    
+    # Clean up the response
+    response_text = response_text.strip()
+    
+    # Remove artifacts and clean up
+    if response_text.startswith("adding"):
+        response_text = "Adding " + response_text[6:]
+    
+    # Take only the first sentence or meaningful phrase
+    sentences = response_text.split('.')
+    if sentences and len(sentences[0].strip()) > 10:
+        review = sentences[0].strip() + "."
+    else:
+        # Fallback to first meaningful line
+        lines = [line.strip() for line in response_text.split('\n') if line.strip()]
+        if lines and len(lines[0]) > 5:
+            review = lines[0]
+            if not review.endswith('.'):
+                review += "."
+        else:
+            review = "Consider adding proper documentation and error handling."
+    
     return review
 
 def parse_ai_response(response_text: str) -> list[ReviewComment]:
@@ -190,7 +205,7 @@ def root_html():
         <ul>
             <li>Mode: AI-Powered</li>
             <li>AI Model: Salesforce/codegen-350M-mono</li>
-            <li>Response Time: ~30-35 seconds</li>
+            <li>Response Time: ~15-25 seconds</li>
         </ul>
     </body>
     </html>