update logic based on official example

app.py

@@ -1,40 +1,45 @@
-from fastapi import FastAPI, HTTPException
+from fastapi import FastAPI, Request
 from pydantic import BaseModel
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
-
 import torch
 
+# Khởi tạo FastAPI app
 app = FastAPI()
 
-# Load model
+# Tải model và tokenizer
 model_name = "VietAI/vit5-base"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
 
-# Input format
-class TextInput(BaseModel):
+# 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):
     text: str
 
 @app.get("/")
-def read_root():
-    return {"message": "ViT5 summarization API is running!"}
+async def root():
+    return {"message": "VietAI vit5-base summarization API is running."}
 
 @app.post("/summarize")
-def summarize(input: TextInput):
-    try:
-        input_text = f"summarize: {input.text}"
-        inputs = tokenizer.encode(input_text, return_tensors="pt", max_length=512, truncation=True)
-        summary_ids = model.generate(
-            inputs,
-            max_length=128,
-            min_length=20,
-            num_beams=4,
-            no_repeat_ngram_size=3,
-            repetition_penalty=2.5,
-            length_penalty=1.0,
-            early_stopping=True
-        )
-        output = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
-        return {"summary": output}
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
+async def summarize(input: SummarizeInput):
+    prefix = "vietnews: "
+    text = prefix + input.text.strip() + " </s>"
+
+    # Tokenize và chuyển sang device
+    encoding = tokenizer(text, return_tensors="pt", max_length=512, truncation=True)
+    input_ids = encoding["input_ids"].to(device)
+    attention_mask = encoding["attention_mask"].to(device)
+
+    # Sinh tóm tắt
+    summary_ids = model.generate(
+        input_ids=input_ids,
+        attention_mask=attention_mask,
+        max_length=256,
+        early_stopping=True
+    )
+
+    summary = tokenizer.decode(summary_ids[0], skip_special_tokens=True, clean_up_tokenization_spaces=True)
+    return {"summary": summary}