File size: 1,808 Bytes
6f98049
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
6fa7da1
6f98049
 
 
 
 
6fa7da1
6f98049
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
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"])