import gradio as gr
import re
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

MODEL_NAME = "csebuetnlp/mT5_multilingual_XLSum"

WHITESPACE_HANDLER = lambda k: re.sub('\s+', ' ', re.sub('\n+', ' ', k.strip()))

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME)

def summarize(text):
    text = WHITESPACE_HANDLER(text)
    
    inputs = tokenizer(text, 
                max_length=1024, 
                truncation=True,
                return_tensors="pt")

    summary_ids = model.generate(inputs["input_ids"])
    summary = tokenizer.batch_decode(summary_ids, 
                                    skip_special_tokens=True, 
                                    clean_up_tokenization_spaces=False)
    return summary


demo = gr.Interface(fn = summarize,
                        inputs = [gr.inputs.Textbox(lines=10, 
                                                    placeholder="Inpuy something...",
                                                    label='Text here !!')],
                        outputs = [gr.outputs.Textbox(
                                                      label="Summary")],
                         
                        title = "🎈 Summarizer 🎈",
                        enable_queue=True)

demo.launch()