# Code 1 - Let's create an UI using Gradio
import gradio as gr
import torch
from transformers import pipeline 

# Code 2 - create a summarizer object
# use torch dtype bfloat16 to compress the model.
summarizer = pipeline(task="summarization", model="facebook/bart-large-cnn", torch_dtype=torch.bfloat16)

# Code 5 - define a function to summarize text
def nlp(input_text):
    summary = summarizer(
        input_text,
        repetition_penalty=5.0,  # Increase this to discourage repetition
        length_penalty=0.3,      # Decrease this to generate longer summaries
        min_length=20, max_length=100
    )
    return summary[0]["summary_text"]


# Code 6 - UI object
ui = gr.Interface(nlp,    
    inputs=gr.Textbox(label="Input Text"), 
    outputs=gr.Textbox(label="Summary"),
    title="Text Summarizer",
    description="Summarize your text using the BART model.")

# Code 7 - launch UI
ui.launch()