add time log and reduce processing time

app.py

@@ -2,42 +2,53 @@ from fastapi import FastAPI, Request
 from pydantic import BaseModel
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 import torch
+from datetime import datetime
+import time
 
 app = FastAPI()
 
-# Load model và tokenizer
+# Load model and tokenizer
 model_name = "VietAI/vit5-base-vietnews-summarization"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-model.to(device)
+model = model.to(device)
 
-# Định nghĩa schema đầu vào
-class SummaryRequest(BaseModel):
+class TextInput(BaseModel):
     text: str
 
 @app.get("/")
 def read_root():
-    return {"message": "VietAI viT5 summarization API is running."}
+    return {"message": "Summarization API is running"}
 
 @app.post("/summarize")
-def summarize(request: SummaryRequest):
-    text = request.text.strip()
-    if not text:
-        return {"summary": ""}
+async def summarize(input_text: TextInput, request: Request):
+    start_time = time.time()
+    print(f"[{datetime.now()}] 🔵 Received request from {request.client.host}")
 
-    prefix = "vietnews: " + text + " </s>"
-    encoding = tokenizer(prefix, return_tensors="pt", truncation=True, max_length=512)
+    text = input_text.text.strip()
+    prefix = "vietnews: "
+    input_text_prefixed = prefix + text + " </s>"
+
+    # Tokenize
+    encoding = tokenizer(input_text_prefixed, return_tensors="pt", truncation=True, max_length=512)
     input_ids = encoding["input_ids"].to(device)
     attention_mask = encoding["attention_mask"].to(device)
 
-    outputs = model.generate(
-        input_ids=input_ids,
-        attention_mask=attention_mask,
-        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)
-    )
+    # Generate summary with optimized settings
+    with torch.inference_mode():
+        outputs = model.generate(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            max_length=96,             # giảm độ dài để xử lý nhanh hơn
+            num_beams=1,               # dùng greedy decoding
+            no_repeat_ngram_size=2,
+            early_stopping=True
+        )
 
     summary = tokenizer.decode(outputs[0], skip_special_tokens=True, clean_up_tokenization_spaces=True)
+
+    end_time = time.time()
+    print(f"[{datetime.now()}] ✅ Response sent — total time: {end_time - start_time:.2f}s")
+
     return {"summary": summary}