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)}"}