Update app.py thirdly

app.py

@@ -26,6 +26,7 @@ model.eval()
 
 # ==== REQUEST / RESPONSE MODELS ====
 
+
 class GenerateRequest(BaseModel):
     prompt: str                     # mô tả cần sinh code
     language: str | None = "Python"
@@ -47,8 +48,8 @@ class CompleteRequest(BaseModel):
     suffix: str = ""                # code phía sau con trỏ (nếu có)
     language: str | None = "Python"
     max_new_tokens: int = 64        # completion thường ngắn
-    num_beams: int = 4
-    temperature: float = 0.7
+    num_beams: int = 1              # completion kiểu Cursor thường để 1 beam
+    temperature: float = 0.3        # ổn định hơn
 
 
 class CodeResponse(BaseModel):
@@ -57,6 +58,7 @@ class CodeResponse(BaseModel):
 
 # ==== TIỆN ÍCH DÙNG CHUNG ====
 
+
 def run_model(prompt: str,
               max_new_tokens: int,
               num_beams: int,
@@ -70,24 +72,32 @@ def run_model(prompt: str,
             num_beams=num_beams,
             temperature=temperature,
             early_stopping=True,
+            repetition_penalty=1.05,   # nhẹ để giảm lặp
         )
 
     text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return text
+    return text.strip()
 
 
 # ==== ENDPOINT 1: TẠO CODE TỪ PROMPT ====
 
+
 @app.post("/generate-code", response_model=CodeResponse)
 def generate_code(req: GenerateRequest):
     lang = req.language or "Python"
 
-    prompt = (
-        f"Generate {lang} code ONLY.\n"
-        f"Do NOT use any other programming language.\n\n"
-        f"Task:\n{req.prompt}\n\n"
-        f"{lang} code:\n"
-    )
+    prompt = f"""
+You are a helpful coding assistant.
+
+Generate ONLY valid {lang} source code for the task below.
+Do NOT add any explanations, comments in natural language, or markdown.
+Return only raw {lang} code.
+
+Task:
+{req.prompt}
+
+{lang} code:
+""".strip()
 
     output = run_model(
         prompt,
@@ -101,16 +111,21 @@ def generate_code(req: GenerateRequest):
 
 # ==== ENDPOINT 2: SỬA LỖI CODE ====
 
+
 @app.post("/fix-code", response_model=CodeResponse)
 def fix_code(req: FixRequest):
     lang = req.language or "Python"
 
-    prompt = (
-        f"The following {lang} code contains bugs.\n"
-        f"Fix all bugs and return ONLY the corrected {lang} code.\n\n"
-        f"Buggy {lang} code:\n{req.code}\n\n"
-        f"Corrected {lang} code:\n"
-    )
+    prompt = f"""
+The following {lang} code contains bugs.
+Fix all bugs and return ONLY the corrected {lang} code.
+Do NOT add any explanations or comments in natural language.
+
+Buggy {lang} code:
+{req.code}
+
+Corrected {lang} code:
+""".strip()
 
     output = run_model(
         prompt,
@@ -122,20 +137,36 @@ def fix_code(req: FixRequest):
     return CodeResponse(output=output)
 
 
-# ==== ENDPOINT 3: GỢI Ý CODE (COMPLETION) ====
+# ==== ENDPOINT 3: GỢI Ý CODE KIỂU CURSOR (COMPLETION) ====
+
 
 @app.post("/complete-code", response_model=CodeResponse)
 def complete_code(req: CompleteRequest):
     lang = req.language or "Python"
 
-    # prefix + suffix giống kiểu Copilot completion
-    prompt = (
-        f"Complete the following {lang} code.\n"
-        f"Only generate the missing code between the prefix and suffix.\n\n"
-        f"Prefix:\n{req.prefix}\n\n"
-        f"Suffix:\n{req.suffix}\n\n"
-        f"Missing {lang} code:\n"
-    )
+    prompt = f"""
+You are an AI code completion engine like Cursor or GitHub Copilot.
+
+You will be given the prefix and suffix of a {lang} file.
+Your task is to generate ONLY the missing {lang} code between them.
+
+Rules:
+- DO NOT repeat the prefix.
+- DO NOT repeat the suffix.
+- DO NOT add any explanations, natural language text, or markdown.
+- DO NOT add imports/includes if they already appear in the prefix.
+- Return ONLY raw {lang} code that can be directly inserted at the cursor.
+
+Prefix:
+{req.prefix}
+
+<CURSOR HERE>
+
+Suffix:
+{req.suffix}
+
+Missing {lang} code:
+""".strip()
 
     output = run_model(
         prompt,
@@ -149,6 +180,7 @@ def complete_code(req: CompleteRequest):
 
 # ==== HEALTHCHECK ====
 
+
 @app.get("/")
 def root():
     return {"status": "CodeT5+ backend is running 🚀"}