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


st.set_page_config(
    page_title="", 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("snoop2head/KoBrailleT5-small-v1")
model = load_model("snoop2head/KoBrailleT5-small-v1")


st.title("한국어 점역과 역점역")
st.write("Braille Pattern Conversion")


default_value = '⠍⠗⠠⠪⠋⠕⠀⠘⠪⠐⠗⠒⠊⠕⠐⠀⠘⠮⠐⠍⠨⠟⠀⠚⠣⠕⠚⠕⠂'
src_text = st.text_area(
    "번역하고 싶은 문장을 입력하세요:",
    default_value,
    height=300,
    max_chars=100,
)
print(src_text)



if src_text == "":
    st.warning("Please **enter text** for translation")
else: 
    # translate into english sentence
    src_text += "</s>"
    translation_result = model.generate(
        tokenizer(
            src_text,
            return_tensors="pt",
            padding="max_length",
            truncation=True,
            max_length=64,
        ).input_ids,
    )
    translation_result = tokenizer.decode(
        translation_result[0],
        clean_up_tokenization_spaces=True,
        skip_special_tokens=True,
    )

    print(f"{src_text} -> {translation_result}")

    st.write(translation_result)
    print(translation_result)