import streamlit as st
from transformers import  AutoModelForSeq2SeqLM, MT5Tokenizer

st.set_page_config(layout="wide", page_title="MT5 Persian Summary")


# Load model lazily
@st.cache(allow_output_mutation=True)
def load_model():
    model = AutoModelForSeq2SeqLM.from_pretrained('nafisehNik/mt5-persian-summary')
    return model

# Load tokenizer lazily
@st.cache(allow_output_mutation=True)
def load_tokenizer():
    tokenizer = MT5Tokenizer.from_pretrained("nafisehNik/mt5-persian-summary")
    return tokenizer

with st.spinner(text="Please wait while the model is loading...."):
    # cache
    fine = load_model()
    tokenizer = load_tokenizer()


# method for summary generation, using the global model and tokenizer
def generate_summary(model, abstract, num_beams = 2, repetition_penalty = 1.0,
                    length_penalty = 2.0, early_stopping = True, max_output_length = 120):
    source_encoding=tokenizer(abstract, max_length=1000, padding="max_length", truncation=True, return_attention_mask=True, add_special_tokens=True, return_tensors="pt")

    generated_ids=model.generate(
        input_ids=source_encoding["input_ids"],
        attention_mask=source_encoding["attention_mask"],
        num_beams=num_beams,
        max_length=max_output_length,
        repetition_penalty=repetition_penalty,
        length_penalty=length_penalty,
        early_stopping=early_stopping,
        use_cache=True
        )

    preds=[tokenizer.decode(gen_id, skip_special_tokens=True, clean_up_tokenization_spaces=True) 
         for gen_id in generated_ids]

    return "".join(preds)


 # APP
# set title and subtitle
st.title("Summarization for Persian")
st.markdown("Demo of nafisehNik/mt5-persian-summary")
st.markdown("Paste any persian (farsi) text you want to summarize.")
# create the input text box and setting panel
num_beams = st.slider("num_beams", min_value=1, max_value=10, value=2)
max_output_length = st.slider("max_output_length", min_value=20, max_value=256, value=120)


text = st.text_area('Paste your text here...','''به گزارش شانا، علی کاردر امروز (۲۷ دی ماه) در مراسم تودیع محسن قمصری، مدیر سابق امور بین الملل شرکت ملی نفت ایران و معارفه سعید خوشرو، مدیر جدید امور بین الملل این شرکت، گفت: مدیریت امور بین الملل به عنوان یکی از تاثیرگذارترین مدیریت های شرکت ملی نفت ایران در دوران تحریم های ظالمانه غرب علیه کشورمان بسیار هوشمندانه عمل کرد و ما توانستیم به خوبی از عهده تحریم ها برآییم. وی افزود: مجموعه امور بین الملل در همه دوران ها با سختی ها و مشکلات بسیاری مواجه بوده است، به ویژه در دوره اخیر به دلیل مسایل پیرامون تحریم وظیفه سنگینی بر عهده داشت که با تدبیر مدیریت خوب این مجموعه سربلند از آن بیرون آمد. کاردر با قدردانی از زحمات محسن قمصری، به سلامت مدیریت امور بین الملل این شرکت اشاره کرد و افزود: محوریت کار مدیریت اموربین الملل سلامت مالی بوده است. وی بر ضرورت نهادینه سازی جوانگرایی در مدیریت شرکت ملی نفت ایران تاکید کرد و گفت: مدیریت امور بین الملل در پرورش نیروهای زبده و کارآزموده آنچنان قوی عملکرده است که برای انتخاب مدیر جدید مشکلی وجود نداشت. کاردر، حرفه ای گری و کار استاندارد را از ویژگی های مدیران این مدیریت برشمرد و گفت: نگاه جامع، خلاقیت و نوآوری و بکارگیری نیروهای جوان باید همچنان مد نظر مدیریت جدید امور بین الملل شرکت ملی نفت ایران باشد.''' ,height=200)
button = st.button("Summarize")

# if button is clicked
with st.spinner("Processing Text  and Summarizing..."):
    if button and text:
        # extract markers from the text
        result = generate_summary(model=fine, abstract=text, num_beams=num_beams, max_output_length=max_output_length)

        # display the result
        st.markdown("**Summary:**")
        st.write(result)