'''
A script that uses gradio to make a web app that takes in a dialog and outputs a summary.
Import a huggingface model.
'''


import gradio as gr
from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM

MODEL_NAME = "Firefly777a/bart-large-samsum-candice-set-summary-3ep-3"

# Import the model
summarizer = pipeline("summarization", model=MODEL_NAME, tokenizer=MODEL_NAME)
# tokenizer = AutoTokenizer.from_pretrained(PRETRAIN_MODEL_NAME)

# Define the function that will be used by the web app
def summarize_dialog(dialog, max_length):
    # Summarize the dialog
    summary = summarizer(dialog, max_length=max_length, min_length=10, do_sample=False)
    # Return the summary
    return summary[0]['summary_text']

# Define the web app & launch
gr.Interface(fn=summarize_dialog, 
             inputs=["text", 
                        gr.Slider(32, 256, value=128)
                        ],
             outputs="text").launch()