vichter's picture
Update app.py
6fa7da1 verified
from fastapi import FastAPI, HTTPException, Header, Depends, Request
from pydantic import BaseModel
from transformers import pipeline
from slowapi import Limiter, _rate_limit_exceeded_handler
from slowapi.util import get_remote_address
from slowapi.errors import RateLimitExceeded
import logging
import os
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
limiter = Limiter(key_func=get_remote_address)
app = FastAPI(title="Panoptifi Summary API")
app.state.limiter = limiter
app.add_exception_handler(RateLimitExceeded, _rate_limit_exceeded_handler)
API_KEY = os.environ.get("API_KEY", "")
def verify_api_key(x_api_key: str = Header(None, alias="X-API-Key")):
if API_KEY and x_api_key != API_KEY:
raise HTTPException(status_code=401, detail="Invalid API key")
return True
logger.info("Loading summarization model...")
summarizer = pipeline(
"summarization",
model="Falconsai/text_summarization"
)
logger.info("Model loaded")
class SummaryInput(BaseModel):
text: str
max_length: int = 150
min_length: int = 30
class SummaryResult(BaseModel):
summary: str
@app.get("/health")
# @limiter.limit("60/minute")
def health(request: Request):
return {"status": "healthy", "model": "Falconsai/text_summarization"}
@app.post("/summarize", response_model=SummaryResult)
# @limiter.limit("20/minute")
def summarize(request: Request, input: SummaryInput, _: bool = Depends(verify_api_key)):
if not input.text.strip():
raise HTTPException(400, "Text cannot be empty")
if len(input.text) < 100:
return SummaryResult(summary=input.text)
result = summarizer(
input.text[:4000],
max_length=input.max_length,
min_length=input.min_length,
do_sample=False
)[0]
return SummaryResult(summary=result["summary_text"])