Update app.py seccondly

app.py

@@ -6,16 +6,15 @@ import torch
 app = FastAPI(title="CodeT5+ Backend on HuggingFace")
 
 # ==== LOAD MODEL ====
-base_ckpt = "Salesforce/codet5p-770m"
-finetuned_ckpt = "Salesforce/codet5p-770m"
+model_name = "Salesforce/codet5p-770m"  # model đa ngôn ngữ, không fine-tune Python-only
 
 print("Loading tokenizer + config...")
-tokenizer = AutoTokenizer.from_pretrained(base_ckpt)
-config = AutoConfig.from_pretrained(base_ckpt)
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+config = AutoConfig.from_pretrained(model_name)
 
-print("Loading fine-tuned model weights...")
+print("Loading model weights...")
 model = T5ForConditionalGeneration.from_pretrained(
-    finetuned_ckpt,
+    model_name,
     config=config
 )
 
@@ -26,45 +25,129 @@ model = model.to(device)
 model.eval()
 
 # ==== REQUEST / RESPONSE MODELS ====
+
 class GenerateRequest(BaseModel):
-    prompt: str
+    prompt: str                     # mô tả cần sinh code
     language: str | None = "Python"
-    task: str = "generate"
     max_new_tokens: int = 128
     num_beams: int = 4
     temperature: float = 0.7
 
-class GenerateResponse(BaseModel):
-    output: str
+
+class FixRequest(BaseModel):
+    code: str                       # code bị lỗi
+    language: str | None = "Python"
+    max_new_tokens: int = 128
+    num_beams: int = 4
+    temperature: float = 0.3        # thấp để sửa lỗi ổn định hơn
+
+
+class CompleteRequest(BaseModel):
+    prefix: str                     # code phía trước con trỏ
+    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
 
 
-def build_prompt(req: GenerateRequest):
-    if req.task == "generate":
-        return f"Generate {req.language} code:\n{req.prompt}"
-    elif req.task == "fix":
-        return f"Fix the bug in the following {req.language} code:\n{req.prompt}\n\nCorrected code:"
-    else:
-        return req.prompt
+class CodeResponse(BaseModel):
+    output: str
 
 
-@app.post("/generate", response_model=GenerateResponse)
-def generate(req: GenerateRequest):
+# ==== TIỆN ÍCH DÙNG CHUNG ====
 
-    prompt = build_prompt(req)
+def run_model(prompt: str,
+              max_new_tokens: int,
+              num_beams: int,
+              temperature: float) -> str:
     inputs = tokenizer(prompt, return_tensors="pt").to(device)
 
     with torch.no_grad():
         outputs = model.generate(
             **inputs,
-            max_new_tokens=req.max_new_tokens,
-            num_beams=req.num_beams,
-            temperature=req.temperature,
-            early_stopping=True
+            max_new_tokens=max_new_tokens,
+            num_beams=num_beams,
+            temperature=temperature,
+            early_stopping=True,
         )
 
     text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return GenerateResponse(output=text)
+    return text
+
+
+# ==== 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"
+    )
+
+    output = run_model(
+        prompt,
+        max_new_tokens=req.max_new_tokens,
+        num_beams=req.num_beams,
+        temperature=req.temperature,
+    )
+
+    return CodeResponse(output=output)
+
+
+# ==== 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"
+    )
+
+    output = run_model(
+        prompt,
+        max_new_tokens=req.max_new_tokens,
+        num_beams=req.num_beams,
+        temperature=req.temperature,
+    )
+
+    return CodeResponse(output=output)
+
+
+# ==== ENDPOINT 3: GỢI Ý CODE (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"
+    )
+
+    output = run_model(
+        prompt,
+        max_new_tokens=req.max_new_tokens,
+        num_beams=req.num_beams,
+        temperature=req.temperature,
+    )
+
+    return CodeResponse(output=output)
+
 
+# ==== HEALTHCHECK ====
 
 @app.get("/")
 def root():