Spaces:
Sleeping
Sleeping
File size: 1,407 Bytes
69847a0 c5a0bf8 20688a8 c5a0bf8 c3ffcdd 4814cd0 29e22ca c5a0bf8 20688a8 4814cd0 69847a0 4814cd0 69847a0 4814cd0 c5a0bf8 4814cd0 c3ffcdd 4814cd0 c3ffcdd 4814cd0 c3ffcdd 69847a0 4814cd0 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 |
from fastapi import FastAPI, Request
from pydantic import BaseModel
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import torch
import time
import logging
app = FastAPI()
# 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.to(device)
class InputText(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}")
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()
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!"}
|