import gradio as gr

from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
from normalizer import normalize


def summarize(text):
    model = AutoModelForSeq2SeqLM.from_pretrained(
        "Afsara/cse_buet_bangla_t5")
    tokenizer = AutoTokenizer.from_pretrained(
        "Afsara/cse_buet_bangla_t5", use_fast=False)

    input_ids = tokenizer(normalize(text), return_tensors="pt").input_ids
    generated_tokens = model.generate(input_ids)
    decoded_tokens = tokenizer.batch_decode(generated_tokens)[0]
    return decoded_tokens


gr_interface = gr.Interface(fn=summarize, inputs="text", outputs="text")

gr_interface.launch()