from transformers import pipeline
import gradio as gr

model_name_converter = {"bart_large":"juliosocher/bart-large-cnn-finetuned-scientific-articles",
                        "mt5-small-finetuned-mt5":"jacks392/mt5-small-finetuned-mt5",
                        "facebook": "facebook/bart-large-cnn",
                        "google" : "google/pegasus-xsum"
                        }
def predict(prompt,model_name, max_length):
    if model_name ==None:
        model_name = "google/pegasus-xsum"
    else:
        model_name = model_name_converter[model_name]
    print('la')
    print(model_name)
    print(max_length)
    model = pipeline("summarization",model = model_name)
    summary = model(prompt,max_length)[0]["summary_text"]
    return summary

def extract_model(option):
        if option ==None:
            model_name = "google/pegasus-xsum"
        else:
            model_name = model_name_converter[option]
             
        return print(model_name)
        
options_1 = model_name_converter.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 = 200, label="the max number of characher for summerized")
    gr.Interface(fn=predict, inputs=[textbox, drop_down, length], outputs = "text")

demo.launch()