import gradio as gr
import spaces
from transformers import pipeline

# Initialize Model
get_completion = pipeline("summarization", model="sshleifer/distilbart-cnn-12-6", device=0)


# @spaces.GPU(duration=120)
# def summarize(input):
#     output = get_completion(input)
#     return output[0]['summary_text']

@spaces.GPU(duration=120)
def summarize(input: str) -> str:
    """
    Summarize the given input text using the sshleifer/distilbart-cnn-12-6 model.

    Args:
        input (str): The text to be summarized.

    Returns:
        str: The summarized version of the input text.
    """
    output: List[Dict[str, str]] = get_completion(input)
    return output[0]['summary_text']


gr.close_all()


####### GRADIO APP #######
title = """<h1 id="title"> Text Summarization </h1>"""

description = """
Summarize any text using the `sshleifer/distilbart-cnn-12-6` model under the hood

- The model used for Summarizing Text [DISTILBART-12-6-CNN](https://huggingface.co/sshleifer/distilbart-cnn-12-6).
"""

css = '''
h1#title {
  text-align: center;
}
'''

theme = gr.themes.Soft()
demo = gr.Blocks(css=css, theme=theme)

with demo:
  gr.Markdown(title)
  gr.Markdown(description)
  interface = gr.Interface(fn=summarize, 
                    inputs=[gr.Textbox(label="Text to Summarize", lines=6)], 
                    outputs=[gr.Textbox(label="Result", lines=3)]
                   )
demo.launch()