summarizer / app.py
spacesedan's picture
feat: returning post_id
20dbd9d
raw
history blame
871 Bytes
from fastapi import FastAPI
from pydantic import BaseModel
from transformers import pipeline
app = FastAPI()
summarizer = pipeline("summarization", model="facebook/bart-large-cnn")
class SummarizationRequest(BaseModel):
inputs: str
post_id: str
class SummarizationResponse(BaseModel):
post_id: str
summary: str
@app.post("/summarize", response_model=SummarizationResponse)
async def summarize_text(request: SummarizationRequest):
input_length = len(request.inputs.split())
max_length = max(50, int(input_length * 0.5))
min_length = max(20, int(input_length * 0.2))
summary = summarizer(request.inputs, max_length=max_length, min_length=min_length, do_sample=False)
return {"post_id": request.post_id, "summary": summary[0]["summary_text"]}
@app.get("/")
def greet_json():
return {"message": "BART Summarizer API is running"}