import gradio as gr
import torch
from transformers import pipeline

summarizer = pipeline("summarization", "jordiclive/flan-t5-3b-summarizer", torch_dtype=torch.bfloat16)

raw_document = 'You must be 18 years old to live or work in New York State...'
prompt = "Produce an article summary of the following news article:"

def summarize(input):

    output = summarizer(
        f"{prompt} {input}",
        num_beams=5,
        min_length=5,
        no_repeat_ngram_size=3,
        truncation=True,
        max_length=512,
    )

    return output

demo = gr.Interface(fn=summarize, inputs="text", outputs="text", title="Text Summarization", description="flat-t5-3b-summarization")
demo.launch()