Update main.py

main.py

@@ -1,5 +1,6 @@
 import os
 os.environ["TRANSFORMERS_CACHE"] = "/tmp"
+
 from fastapi import FastAPI
 from pydantic import BaseModel
 from transformers import PegasusTokenizer, PegasusForConditionalGeneration
@@ -7,32 +8,44 @@ import torch
 
 app = FastAPI()
 
-# Load mô hình
+# Load model và tokenizer
 model_name = "google/pegasus-cnn_dailymail"
 tokenizer = PegasusTokenizer.from_pretrained(model_name)
 model = PegasusForConditionalGeneration.from_pretrained(model_name)
+
+# Dùng GPU nếu có
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model = model.to(device)
 
-# Schema input
+# Định nghĩa input schema
 class InputText(BaseModel):
     text: str
 
-# Hàm tóm tắt
+# Hàm tóm tắt tự động điều chỉnh độ dài theo số token
 def summarize(text: str) -> str:
+    # Tokenize input text
     inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=1024)
+    input_length = inputs["input_ids"].shape[1]
+
+    # Xác định độ dài summary theo tỷ lệ input
+    summary_max_len = max(30, int(input_length * 0.2))     # tối đa khoảng 20% số token
+    summary_min_len = max(15, int(summary_max_len * 0.6))  # tối thiểu khoảng 60% max
+
     inputs = {k: v.to(device) for k, v in inputs.items()}
+
+    # Sinh summary
     summary_ids = model.generate(
         inputs["input_ids"],
-        max_length=150,
-        min_length=60,
+        max_length=summary_max_len,
+        min_length=summary_min_len,
         num_beams=4,
         no_repeat_ngram_size=3,
         early_stopping=True
     )
+
     return tokenizer.decode(summary_ids[0], skip_special_tokens=True)
 
-# Route API
+# API route
 @app.post("/summarize")
 def summarize_api(input: InputText):
     return {"summary": summarize(input.text)}