app.py

import gradio as gr
from util import update
from util import summarize
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, DataCollatorForSeq2Seq, Seq2SeqTrainer, TrainingArguments

cp_aug = 'minnehwg/finetune-newwiki-summarization-ver-augmented'

tokenizer = AutoTokenizer.from_pretrained("VietAI/vit5-base")
model  = AutoModelForSeq2SeqLM.from_pretrained(cp_aug)

def summarize_text(text):
    text = summarize(text, model, tokenizer, 4)
    return text


with gr.Blocks() as demo:
    gr.Markdown("Start typing below and then click **Run** to see the output.")
    with gr.Row():
        inp = gr.Textbox(placeholder="What is your name?")
        out = gr.Textbox()
    btn = gr.Button("Run")
    btn.click(fn=summarize_text, inputs=inp, outputs=out)

demo.launch()