from transformers import pipeline
import gradio as gr


MODELS = {
    "gsarti": pipeline("summarization", model="gsarti/it5-base-wiki-summarization"),
    "facebook": pipeline("summarization", model="facebook/bart-large-cnn"),
    "lincoln": pipeline(
        "summarization", model="lincoln/mbart-mlsum-automatic-summarization"
    ),
    "t5-small": pipeline("summarization", model="t5-small"),
}


def predict(prompt, model_name, max_length):
    if model_name is None:
        model = MODELS["t5-small"]
    else:
        model = MODELS[model_name]
    prompt = prompt.replace("\n", " ")
    summary = model(prompt, max_length)[0]["summary_text"]
    return summary


options_1 = list(MODELS.keys())
with gr.Blocks() as demo:
    drop_down = gr.Dropdown(choices=options_1, label="model")
    textbox = gr.Textbox(placeholder="Enter text block to summarize", lines=4)
    length = gr.Number(value=100, label="the max number of characher for summerized")
    gr.Interface(fn=predict, inputs=[textbox, drop_down, length], outputs="text")

demo.launch()