# -*- 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"
)

@st.cache
def load_model(model_name):
    model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
    return model

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


st.title("🤖 KoQuillBot")


default_value = "프로젝트 가치가 미화 1백만 달러 이상인 공공 파트너가 시작한 PPP 프로젝트에 대해 2단계 입찰이 실시됩니다. 입찰을 전자 방식으로 진행하는 것이 허용됩니다. (즉, 신청서 및 입찰 제안의 전자 제출). COVID-19 전염병과 그에 따른 여행 제한으로 인해 오늘날에는 일반적인 관행이 되었습니다."
src_text = st.text_area(
    "바꾸고 싶은 문장을 입력하세요:",
    default_value,
    height=300,
    max_chars=200,
)
print(src_text)



if src_text == "":
    st.warning("Please **enter text** for translation")

# translate into english sentence
english_translation = ko2en_model.generate(
    **tokenizer(
        src_text,
        return_tensors="pt",
        padding="max_length",
        truncation=True,
        max_length=64,
    ),
    max_length=64,
    num_beams=5,
    repetition_penalty=1.3,
    no_repeat_ngram_size=3,
    num_return_sequences=1,
)
english_translation = tokenizer.decode(
    english_translation[0],
    clean_up_tokenization_spaces=True,
    skip_special_tokens=True,
)

# translate back to korean
korean_translation = en2ko_model.generate(
    **tokenizer(
        english_translation,
        return_tensors="pt",
        padding="max_length",
        truncation=True,
        max_length=64,
    ),
    max_length=64,
    num_beams=5,
    repetition_penalty=1.3,
    no_repeat_ngram_size=3,
    num_return_sequences=1,
)

korean_translation = tokenizer.decode(
    korean_translation[0],
    clean_up_tokenization_spaces=True,
    skip_special_tokens=True,
)
print(f"{src_text} -> {english_translation} -> {korean_translation}")

st.write(korean_translation)
print(korean_translation)