| 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] | |
| model = pipeline("summarization",model = model_name) | |
| summary = model(prompt,max_length)[0]["summary_text"] | |
| return summary | |
| options_1 = list(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() | |