import gradio as gr
from gradio.components import Textbox
import re
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

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

model_name = "csebuetnlp/mT5_multilingual_XLSum"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)

def predict(text):
  input_ids = tokenizer(
      [WHITESPACE_HANDLER(text)],
      return_tensors="pt",
      padding="max_length",
      truncation=True,
      max_length=512
  )["input_ids"]

  output_ids = model.generate(
      input_ids=input_ids,
      max_length=84,
      no_repeat_ngram_size=2,
      num_beams=4
  )[0]

  summary = tokenizer.decode(
      output_ids,
      skip_special_tokens=True,
      clean_up_tokenization_spaces=False
  )
  return summary


title = "language-summarize-multilang-mT5"
description = "Instructions: Enter text in Text boxes, then click [Submit] to get summary. Alternatively, you can select examples at the bottom, then click [Submit]."

examples = [
    ["寿司は、日本の伝統的な料理であり、主に酢飯（すしめし）と呼ばれる酢で調味されたご飯と、新鮮な魚や貝、野菜などが使われる食べ物です。寿司は約千年以上前の日本の江戸時代に始まり、保存食から発展し、現代ではその繊細で美しい見た目と素材の鮮度や調和が重視される味わいが特徴的です。寿司は日本国内だけでなく、世界中で広く愛され、さまざまな種類やアレンジが存在しています。寿司を食べる際には、醤油や生姜、わさびなどを添えて楽しむことができ、日本の文化や食の象徴として、多くの人々に親しまれています。"],
    ["News: xAI, Elon Musk’s newly formed AI company, has revealed itself with a new website detailing its mission and team at https://x.ai/. Musk tweeted the company’s intent is to understand reality without any other details or explanation. The goal of xAI is to understand the true nature of the universe, according to the website. The team is headed up by Elon Musk and includes team members that have worked at other big names."],
    ["นายฉัตรชัย ศิริไล ผู้จัดการธนาคารเพื่อการเกษตรและสหกรณ์การเกษตร (ธ.ก.ส.) เปิดเผยว่า ธ.ก.ส. ได้มีการดำเนินนโยบายการแก้ไขปัญหาหนี้ครัวเรือนสอดรับกับนโยบายธนาคารแห่งประเทศไทย (ธปท.) โดยแบ่งกลุ่มลูกหนี้ตามศักยภาพเป็น 3 กลุ่ม ได้แก่ กลุ่มลูกหนี้ปกติ กลุ่มลูกหนี้ Hybrid และกลุ่มลูกหนี้ NPL และจัดทำเครื่องมือหรือวิธีการแก้ไขหนี้ให้สอดคล้องกับปัญหาที่แท้จริงของลูกหนี้เฉพาะกลุ่ม ควบคู่กับการเสริมความรู้ ทางการเงิน เทคโนโลยี เพื่อเพิ่มประสิทธิภาพในการประกอบอาชีพและยกระดับรายได้"],
]

demo = gr.Interface(
    predict,
    inputs=[
        Textbox(label="Text",lines = 8),
    ],
    outputs=[Textbox(label="Summary")],
    title=title,
    description=description,
    examples=examples,
    theme="freddyaboulton/dracula_revamped",
    )

demo.launch()