import torch
import gradio as gr

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

text_summary = pipeline("summarization", model="sshleifer/distilbart-cnn-12-3", torch_dtype= torch.bfloat16)

# model_path = "../Models/models--sshleifer--distilbart-cnn-12-6/snapshots/a4f8f3ea906ed274767e9906dbaede7531d660ff"
# text_summary = pipeline("summarization", model=model_path, torch_dtype= torch.bfloat16)

# text = "A computer is a machine that can be programmed to automatically carry out sequences of arithmetic " \
# "or logical operations (computation). Modern digital electronic computers can perform generic sets of " \
# "operations known as programs, which enable computers to perform a wide range of tasks. The term computer " \
# "system may refer to a nominally complete computer that includes the hardware, operating system, software, " \
# "and peripheral equipment needed and used for full operation; or to a group of computers that are linked and " \
# "function together, such as a computer network or computer cluster."
# print(text_summary(text))


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

gr.close_all()

# demo = gr.Interface(fn=summary, inputs="text", outputs="text")
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 1. Text Summarizer",
                    description="This application will be used to summarize the text using Hugging-Face pre-trained model")
demo.launch()