update logic based on official example

app.py

@@ -3,43 +3,41 @@ from pydantic import BaseModel
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 import torch
 
-# Khởi tạo FastAPI app
 app = FastAPI()
 
-# Tải model và tokenizer
+# Load model và tokenizer
 model_name = "VietAI/vit5-base"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
-
-# Thiết bị (GPU nếu có, nếu không dùng CPU)
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model.to(device)
 
-# Schema cho input
-class SummarizeInput(BaseModel):
+# Định nghĩa schema đầu vào
+class SummaryRequest(BaseModel):
     text: str
 
 @app.get("/")
-async def root():
-    return {"message": "VietAI vit5-base summarization API is running."}
+def read_root():
+    return {"message": "VietAI viT5 summarization API is running."}
 
 @app.post("/summarize")
-async def summarize(input: SummarizeInput):
-    prefix = "vietnews: "
-    text = prefix + input.text.strip() + " </s>"
+def summarize(request: SummaryRequest):
+    text = request.text.strip()
+    if not text:
+        return {"summary": ""}
 
-    # Tokenize và chuyển sang device
-    encoding = tokenizer(text, return_tensors="pt", max_length=512, truncation=True)
+    prefix = "vietnews: " + text + " </s>"
+    encoding = tokenizer(prefix, return_tensors="pt", truncation=True, max_length=512)
     input_ids = encoding["input_ids"].to(device)
     attention_mask = encoding["attention_mask"].to(device)
 
-    # Sinh tóm tắt
-    summary_ids = model.generate(
+    outputs = model.generate(
         input_ids=input_ids,
         attention_mask=attention_mask,
-        max_length=256,
-        early_stopping=True
+        max_length=128,     # Tóm tắt ngắn gọn
+        do_sample=False,    # Không sampling
+        num_beams=1         # Greedy decoding (nhanh nhất)
     )
 
-    summary = tokenizer.decode(summary_ids[0], skip_special_tokens=True, clean_up_tokenization_spaces=True)
+    summary = tokenizer.decode(outputs[0], skip_special_tokens=True, clean_up_tokenization_spaces=True)
     return {"summary": summary}