Summary / app.py
Uhhy's picture
Update app.py
8bd94c7 verified
raw
history blame contribute delete
970 Bytes
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from multiprocessing import Process, Queue
from transformers import pipeline
import uvicorn
app = FastAPI()
summarization_pipeline = pipeline("summarization", model="facebook/bart-large-cnn")
class SummaryRequest(BaseModel):
text: str
def generate_summary_response(request, queue):
try:
summary = summarization_pipeline(request.text)
queue.put(summary[0]['summary_text'])
except Exception as e:
queue.put(f"Error: {str(e)}")
@app.post("/generate_summary")
async def generate_summary(request: SummaryRequest):
queue = Queue()
p = Process(target=generate_summary_response, args=(request, queue))
p.start()
p.join()
response = queue.get()
if "Error" in response:
raise HTTPException(status_code=500, detail=response)
return {"summary": response}
if __name__ == "__main__":
uvicorn.run(app, host="0.0.0.0", port=7860)