CazimirRoman's picture
replace summarizer pipeline with a faster version
93795ea
raw
history blame contribute delete
No virus
963 Bytes
from typing import Union
from fastapi import FastAPI
from transformers import pipeline
from newspaper import Article
app = FastAPI()
# see dockerfile. it has a line to download the model when the docket container is initialized
# summarizer = pipeline("summarization", "/code/models")
summarizer = pipeline("summarization", model="Falconsai/text_summarization")
@app.get("/")
def read_root():
return {"Hello": "World"}
@app.get("/items/{item_id}")
def read_item(item_id: int, q: Union[str, None] = None):
return {"item_id": item_id, "q": q}
@app.post("/summarize")
async def summarize(text: str):
# url = 'https://towardsdatascience.com/containerizing-huggingface-transformers-for-gpu-inference-with-docker-and-fastapi-on-aws-d4a83edede2f'
article = Article(text)
article.download()
article.parse()
return {"summary": f"{summarizer(article.text, max_length=130, min_length=30, do_sample=False)}"}