import torch
import gradio as gr

# Use a pipeline as a high-level helper
from transformers import pipeline

# downloaded the model from web
text_summary = pipeline("summarization", model="sshleifer/distilbart-cnn-12-6",
               torch_dtype=torch.bfloat16)

def summary(input):
    output = text_summary(input)
    return output[0]['summary_text']


gr.close_all()

# simple gradio web app
# demo = gr.Interface(fn=summary, inputs="text", outputs="text")

# beautified
demo = gr.Interface(
    fn=summary,
    inputs=[gr.Textbox(label="Input text to summarize", lines=6)],
    outputs=[gr.Textbox(label="Summarized text", lines=4)],
    title="Project 01: Text Summarization",
    description="As understood from the title, if not already, this application will summarize your text"
)

demo.launch()