add time log and reduce processing time

app.py

@@ -2,53 +2,48 @@ from fastapi import FastAPI, Request
 from pydantic import BaseModel
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 import torch
-from datetime import datetime
 import time
+import logging
 
 app = FastAPI()
 
-# Load model and tokenizer
-model_name = "VietAI/vit5-base-vietnews-summarization"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+# Logging setup
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger("summarizer")
+
+# Model & 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 = model.to(device)
+model.to(device)
 
-class TextInput(BaseModel):
+class InputText(BaseModel):
     text: str
 
-@app.get("/")
-def read_root():
-    return {"message": "Summarization API is running"}
-
 @app.post("/summarize")
-async def summarize(input_text: TextInput, request: Request):
+async def summarize(req: Request, input: InputText):
     start_time = time.time()
-    print(f"[{datetime.now()}] 🔵 Received request from {request.client.host}")
-
-    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)
-
-    # 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)
+    logger.info(f"\U0001F535 Received request from {req.client.host}")
+
+    text = input.text.strip()
+    inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=512).to(device)
+
+    outputs = model.generate(
+        **inputs,
+        max_length=128,
+        num_beams=2,
+        no_repeat_ngram_size=2,
+        early_stopping=True
+    )
+    summary = tokenizer.decode(outputs[0], skip_special_tokens=True)
 
     end_time = time.time()
-    print(f"[{datetime.now()}] ✅ Response sent — total time: {end_time - start_time:.2f}s")
+    duration = end_time - start_time
+    logger.info(f"\u2705 Response sent — total time: {duration:.2f}s")
 
     return {"summary": summary}
+
+@app.get("/")
+def root():
+    return {"message": "Vietnamese Summarization API is up and running!"}