import streamlit as st
from transformers import (AutoModel, AutoTokenizer, pipeline, T5ForConditionalGeneration)

# import torch

st.set_page_config(layout="wide", page_title="تبدیل متون محاوره‌ای به فارسی")

model_path = 'erfan226/persian-t5-formality-transfer'
model = T5ForConditionalGeneration.from_pretrained(model_path)
tokenizer = AutoTokenizer.from_pretrained(model_path)
pipe = pipeline(task='text2text-generation', model=model, tokenizer=tokenizer)

def paraphrase(text, min_len, max_len=128, n_beams=5):
  for j in range(3):
    out = pipe(text, encoder_no_repeat_ngram_size=4, do_sample=True, num_beams=n_beams, max_length=max_len)[0]['generated_text']
    print("Paraphrase:", out)

text = st.text_area('Paste your link here...', "من با دوستام میرم بازی", height=50)
num_beams = st.sidebar.slider('طول Beam', min_value=1, max_value=10, value=5, step=1)
min_length = st.sidebar.slider('حداقل طول جمله', min_value=10, max_value=100, value=50, step=10)
max_length = st.sidebar.slider('حداکثر طول طول جمله', min_value=30, max_value=700, value=100, step=10)
button = st.button("تبدیل متن")

with st.spinner("در حال تبدیل متن..."):
    # if button and text:
    # get the text
    out = paraphrase(text, num_beams, min_length, max_length)
    st.write(out)
    # display the summary
    st.markdown("**متن خروجی:**")
    st.write(out)

# x = st.slider('Select a value')
# st.write(x, 'squared is', x * x)