# -*- coding: utf-8 -*-
import numpy as np
import streamlit as st
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM


st.set_page_config(page_title='KoQuillBot', layout='wide', initial_sidebar_state='expanded')

tokenizer = AutoTokenizer.from_pretrained("QuoQA-NLP/KE-T5-Ko2En-Base")
ko2en_model = AutoModelForSeq2SeqLM.from_pretrained("QuoQA-NLP/KE-T5-Ko2En-Base")
en2ko_model = AutoModelForSeq2SeqLM.from_pretrained("QuoQA-NLP/KE-T5-En2Ko-Base")


st.title("🤖 KoQuillBot")

src_text = st.text_area("바꾸고 싶은 문장을 입력하세요:",height=None,max_chars=None,key=None,help="Enter your text here")

backtranslated = ""

if st.button('문장 변환'):
    if src_text == "":
        st.warning('Please **enter text** for translation')

    else:
        translated = ko2en_model.generate(
        **tokenizer([src_text], return_tensors="pt", padding=True, max_length=64,),
        max_length=64,
        num_beams=5,
        repetition_penalty=1.3,
        no_repeat_ngram_size=3,
        num_return_sequences=1,
        ) 
        
        backtranslated = en2ko_model.generate(
        **tokenizer([translated], return_tensors="pt", padding=True, max_length=64,),
        max_length=64,
        num_beams=5,
        repetition_penalty=1.3,
        no_repeat_ngram_size=3,
        num_return_sequences=1,
        )
else:
    pass


print(backtranslated)

st.write(backtranslated)