app.py

from fastapi import FastAPI, Request
from pydantic import BaseModel
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import torch

# Khởi tạo FastAPI app
app = FastAPI()

# Tải model và tokenizer
model_name = "VietAI/vit5-base"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)

# 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("/")
async def root():
    return {"message": "VietAI vit5-base summarization API is running."}

@app.post("/summarize")
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}