add time log and reduce processing time

app.py

@@ -1,49 +1,65 @@
+import time
+import logging
+import torch
 from fastapi import FastAPI, Request
 from pydantic import BaseModel
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
-import torch
-import time
-import logging
+from concurrent.futures import ThreadPoolExecutor
+import asyncio
 
+# Khởi tạo app
 app = FastAPI()
 
-# Logging setup
+# Logging
 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)
+# Load model và tokenizer
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-model.to(device)
+tokenizer = AutoTokenizer.from_pretrained("VietAI/vit5-base")
+model = AutoModelForSeq2SeqLM.from_pretrained("VietAI/vit5-base").to(device)
 
-class InputText(BaseModel):
+# Thread executor để xử lý blocking
+executor = ThreadPoolExecutor(max_workers=2)
+
+# Kiểu dữ liệu đầu vào
+class TextIn(BaseModel):
     text: str
 
-@app.post("/summarize")
-async def summarize(req: Request, input: InputText):
-    start_time = time.time()
-    logger.info(f"\U0001F535 Received request from {req.client.host}")
+# -------------------------------
+# GET: kiểm tra API sẵn sàng
+@app.get("/")
+def read_root():
+    return {"message": "API is ready."}
 
-    text = input.text.strip()
-    inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=512).to(device)
+# -------------------------------
+# Hàm tóm tắt (blocking)
+def summarize_text(text: str) -> str:
+    prompt = "vietnews: " + text.strip() + " </s>"
+    encoding = tokenizer(prompt, 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(
-        **inputs,
+        input_ids=input_ids,
+        attention_mask=attention_mask,
         max_length=128,
         num_beams=2,
-        no_repeat_ngram_size=2,
         early_stopping=True
     )
-    summary = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return tokenizer.decode(outputs[0], skip_special_tokens=True, clean_up_tokenization_spaces=True)
+
+# -------------------------------
+# POST: async API tóm tắt
+@app.post("/summarize")
+async def summarize(request: Request, payload: TextIn):
+    start_time = time.time()
+    client_ip = request.client.host
+    logging.info(f"[{time.strftime('%Y-%m-%d %H:%M:%S')}] 🔵 Received request from {client_ip}")
+
+    summary = await asyncio.get_event_loop().run_in_executor(executor, summarize_text, payload.text)
 
     end_time = time.time()
     duration = end_time - start_time
-    logger.info(f"\u2705 Response sent — total time: {duration:.2f}s")
+    logging.info(f"[{time.strftime('%Y-%m-%d %H:%M:%S')}] ✅ Response sent — total time: {duration:.2f}s")
 
     return {"summary": summary}
-
-@app.get("/")
-def root():
-    return {"message": "Vietnamese Summarization API is up and running!"}